# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

%pip install --upgrade --quiet google-cloud-speech ipywebrtc

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

import json
import os
import re
from collections.abc import Generator

from IPython.display import HTML, Audio, display
from google.api_core.client_options import ClientOptions
from google.cloud.speech_v2 import SpeechClient
from google.cloud.speech_v2.types import cloud_speech
from ipywebrtc import AudioRecorder, CameraStream

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

STT_LOCATION = "us"  # @param {type: "string"}

! gcloud config set project {PROJECT_ID}
! gcloud auth application-default login -q
! gcloud auth application-default set-quota-project {PROJECT_ID}

client = SpeechClient(
    client_options=ClientOptions(api_endpoint=f"{STT_LOCATION}-speech.googleapis.com")
)
recognizer = client.recognizer_path(PROJECT_ID, STT_LOCATION, "_")
model = "chirp_3"
# Set a timeout for the batch recognition operation
MAX_AUDIO_LENGTH_SECS = 8 * 60 * 60

def print_transcript(results: list[cloud_speech.SpeechRecognitionResult]) -> None:
    for result in results:
        display(
            HTML(
                f"""<div style="word-break: break-all;">{result.alternatives[0].transcript}</div>"""
            )
        )


def generate_audio_chunks(
    audio_content: bytes, chunk_size: int
) -> Generator[bytes, None, None]:
    """Splits a byte string (like audio data) into smaller, equal-sized chunks.

    Args:
        audio_content: The raw byte data of the audio.
        chunk_size: The desired size for each audio chunk in bytes.

    Yields:
        A series of audio chunks.
    """
    # Loop through the audio content, stepping by chunk_size each time
    for start_index in range(0, len(audio_content), chunk_size):
        # The end_index is the start_index plus the chunk size
        end_index = start_index + chunk_size
        # Yield the slice of audio data for the current chunk
        yield audio_content[start_index:end_index]


def group_utterances_by_speaker_from_file(json_file_path: str) -> dict:
    """Reads a JSON file containing transcribed words and groups them into sentences spoken by each speaker."""
    with open(json_file_path, encoding="utf-8") as f:
        json_data_string = f.read()
    words_regex = r'"words":\s*(\[.*?\])'
    match = re.search(words_regex, json_data_string, re.DOTALL)

    words_list = json.loads(match.group(1))
    dialogue = []
    current_speaker = None
    current_utterance_words = []
    current_speaker = words_list[0]["speakerLabel"]

    for item in words_list:
        word = item["word"]
        speaker = item["speakerLabel"]
        # Check if the speaker has changed
        if speaker != current_speaker:
            dialogue.append(
                {"speaker": current_speaker, "text": " ".join(current_utterance_words)}
            )
            # Start a new utterance
            current_speaker = speaker
            current_utterance_words = [word]
        else:
            # Continue the current utterance
            current_utterance_words.append(word)
    # Add the final pending utterance
    if current_speaker is not None:
        dialogue.append(
            {"speaker": current_speaker, "text": " ".join(current_utterance_words)}
        )

    return {"dialogue": dialogue}

audio_url = "https://storage.googleapis.com/cloud-samples-data/generative-ai/audio/audio_summary_clean_energy_short.mp3"
audio_filename = os.path.basename(audio_url)
! wget {audio_url} -O {audio_filename}

display(Audio(filename=audio_filename))

config = cloud_speech.RecognitionConfig(
    auto_decoding_config=cloud_speech.AutoDetectDecodingConfig(),
    model=model,
    language_codes=["en-US"],
)

with open(audio_filename, "rb") as f:
    audio_content = f.read()

request = cloud_speech.RecognizeRequest(
    recognizer=recognizer,
    config=config,
    content=audio_content,
)

response = client.recognize(request=request)

print_transcript(response.results)

audio_url = (
    "https://storage.googleapis.com/cloud-samples-data/generative-ai/audio/spanish.wav"
)
audio_gcs_uri = audio_url.replace("https://storage.googleapis.com/", "gs://")

display(Audio(url=audio_url))

config = cloud_speech.RecognitionConfig(
    auto_decoding_config=cloud_speech.AutoDetectDecodingConfig(),
    model=model,
    language_codes=["auto"],
)
request = cloud_speech.RecognizeRequest(
    recognizer=recognizer,
    config=config,
    uri=audio_gcs_uri,
)

response = client.recognize(request=request)

print_transcript(response.results)

audio_url = "https://storage.googleapis.com/cloud-samples-data/generative-ai/audio/Chirp-3-Docs-Dive.mp3"
audio_gcs_uri = audio_url.replace("https://storage.googleapis.com/", "gs://")
display(Audio(url=audio_url))

gcs_output_folder = "gs://[your-bucket-path]"  # @param {type: "string"}

config = cloud_speech.RecognitionConfig(
    auto_decoding_config=cloud_speech.AutoDetectDecodingConfig(),
    features=cloud_speech.RecognitionFeatures(
        diarization_config=cloud_speech.SpeakerDiarizationConfig(),
    ),
    model=model,
    language_codes=["en-US"],
)

files = [cloud_speech.BatchRecognizeFileMetadata(uri=audio_gcs_uri)]

request = cloud_speech.BatchRecognizeRequest(
    recognizer=recognizer,
    config=config,
    files=files,
    recognition_output_config=cloud_speech.RecognitionOutputConfig(
        gcs_output_config=cloud_speech.GcsOutputConfig(uri=gcs_output_folder),
    ),
)
operation = client.batch_recognize(request=request)
response = operation.result(timeout=MAX_AUDIO_LENGTH_SECS)

transcript = response.results[audio_gcs_uri].uri

!gsutil cp {transcript} output.json

print(json.dumps(group_utterances_by_speaker_from_file("output.json"), indent=4))

if "google.colab" in sys.modules:
    from google.colab import output

    output.enable_custom_widget_manager()

camera = CameraStream(constraints={"audio": True, "video": False})
recorder = AudioRecorder(stream=camera)
recorder

with open("recording.webm", "wb") as f:
    f.write(recorder.audio.value)

!ffmpeg -i recording.webm -vn -ar 44100 -ac 2 -f mp3 recording.mp3
audio_file = "recording.mp3"

CHUNK_SIZE = 3200

recognition_config = cloud_speech.RecognitionConfig(
    auto_decoding_config=cloud_speech.AutoDetectDecodingConfig(),
    language_codes=["auto"],
    model=model,
)


def create_streaming_requests(
    audio_file_path: str,
) -> Generator[cloud_speech.StreamingRecognizeRequest, None, None]:
    """Prepares and yields all necessary requests for streaming speech recognition.

    First, it yields the configuration request, then it reads an audio file
    and yields its content in chunks.

    Args:
        audio_file_path: The path to the local audio file.

    Yields:
        A stream of StreamingRecognizeRequest objects.
    """
    # a. First, yield the initial configuration request.
    config_request = cloud_speech.StreamingRecognizeRequest(
        recognizer=recognizer,
        streaming_config=cloud_speech.StreamingRecognitionConfig(
            config=recognition_config,
        ),
    )
    yield config_request

    # b. Second, read the audio file and stream it in chunks.
    with open(audio_file_path, "rb") as f:
        audio_content = f.read()

    for chunk in generate_audio_chunks(audio_content, CHUNK_SIZE):
        yield cloud_speech.StreamingRecognizeRequest(audio=chunk)


responses = client.streaming_recognize(requests=create_streaming_requests(audio_file))

print("Streaming transcripts:")
all_transcripts = []
for response in responses:
    for result in response.results:
        transcript = result.alternatives[0].transcript
        print(transcript)
        all_transcripts.append(transcript)

final_transcript = " ".join(all_transcripts)
print(f"\n--- Final Combined Transcript ---\n{final_transcript}")