# 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 gradio

# Use the environment variable if the user doesn't provide Project ID.
import os

# 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"))

TTS_LOCATION = "global"

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

import google.auth
import google.auth.transport.requests

credentials, _ = google.auth.default()
authentication = google.auth.transport.requests.Request()
credentials.refresh(authentication)

import base64
import json
import os

import gradio as gr
import numpy as np
import requests
from IPython.display import Audio, display

API_ENDPOINT = (
    f"{TTS_LOCATION}-texttospeech.googleapis.com"
    if TTS_LOCATION != "global"
    else "texttospeech.googleapis.com"
)

ACCESS_TOKEN = credentials.token

def create_instant_custom_voice_key(
    reference_audio_bytes: bytes, consent_audio_bytes: bytes
) -> str:
    """Creates a temporary custom voice key"""
    url = f"https://{API_ENDPOINT}/v1beta1/voices:generateVoiceCloningKey"

    request_body = {
        "reference_audio": {
            "audio_config": {"audio_encoding": "LINEAR16", "sample_rate_hertz": 24000},
            "content": reference_audio_bytes,
        },
        "voice_talent_consent": {
            "audio_config": {"audio_encoding": "LINEAR16", "sample_rate_hertz": 24000},
            "content": consent_audio_bytes,
        },
        "consent_script": "I am the owner of this voice and I consent to Google using this voice to create a synthetic voice model.",
        "language_code": "en-US",
    }

    response = None
    try:
        headers = {
            "Authorization": f"Bearer {ACCESS_TOKEN}",
            "x-goog-user-project": PROJECT_ID,
            "Content-Type": "application/json; charset=utf-8",
        }

        response = requests.post(url, headers=headers, json=request_body)
        response.raise_for_status()

        response_json = response.json()
        return response_json.get("voiceCloningKey")

    except requests.exceptions.RequestException as e:
        print(f"Error making API request: {e}")
        if response is not None and response.text:
            print("Response error message:")
            print(response.text)
    except json.JSONDecodeError as e:
        print(f"Error decoding JSON response: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")


def synthesize_text_with_cloned_voice(voice_key: str, text: str) -> None:
    """Synthesizes text with the cloned voice"""
    url = f"https://{API_ENDPOINT}/v1beta1/text:synthesize"

    request_body = {
        "input": {"text": text},
        "voice": {
            "language_code": "en-US",
            "voice_clone": {
                "voice_cloning_key": voice_key,
            },
        },
        "audioConfig": {"audioEncoding": "LINEAR16", "sample_rate_hertz": 24000},
    }

    try:
        headers = {
            "Authorization": f"Bearer {ACCESS_TOKEN}",
            "x-goog-user-project": PROJECT_ID,
            "Content-Type": "application/json; charset=utf-8",
        }

        response = requests.post(url, headers=headers, json=request_body)
        response.raise_for_status()

        response_json = response.json()
        audio_content = response_json.get("audioContent")

        if audio_content:
            display(Audio(base64.b64decode(audio_content), rate=24000))
        else:
            print("Error: Audio content not found in the response.")
            print(response_json)

    except requests.exceptions.RequestException as e:
        print(f"Error making API request: {e}")
    except json.JSONDecodeError as e:
        print(f"Error decoding JSON response: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")


def wav_to_base64(file_path: str) -> str:
    """Convert a WAV file to base64 encoded string"""
    try:
        with open(file_path, "rb") as wav_file:
            encoded_string = base64.b64encode(wav_file.read()).decode("utf-8")
            return encoded_string
    except FileNotFoundError:
        print(f"Error: File not found at {file_path}")
        return None
    except Exception as e:
        print(f"An error occurred: {e}")
        return None


def create_voice(
    reference_audio: gr.Audio,
    consent_audio: gr.Audio,
    progress: gr.Progress | None = None,
) -> str:
    """Create a custom voice using reference and consent audio"""
    if reference_audio is None or consent_audio is None:
        return "Please upload both reference and consent audio files."

    if not progress:
        progress = gr.Progress()

    progress(0.2, desc="Processing audio files...")
    reference_audio_b64 = wav_to_base64(reference_audio)
    consent_audio_b64 = wav_to_base64(consent_audio)

    if reference_audio_b64 is None or consent_audio_b64 is None:
        return "Error processing audio files."

    progress(0.5, desc="Creating voice clone...")
    voice_key = create_instant_custom_voice_key(reference_audio_b64, consent_audio_b64)

    if voice_key:
        progress(1.0, desc="Voice created successfully!")
        return voice_key
    return "Failed to create voice. Check the logs for details."


def create_voice_with_masking(
    reference_audio: gr.Audio, consent_audio: gr.Audio
) -> tuple:
    """A wrapper function for create_voice to handle masking"""
    key = create_voice(reference_audio, consent_audio)
    if key and len(key) > 8:
        masked_key = key[:5] + "..."
    else:
        masked_key = key
    return key, masked_key


def generate_speech(
    voice_key: str, text: str, progress: gr.Progress | None = None
) -> tuple:
    """Generate speech using the cloned voice"""
    if not voice_key or not text:
        return None, "Please create a voice key and enter text to synthesize."

    if not progress:
        progress = gr.Progress()

    progress(0.3, desc="Generating speech...")

    try:
        url = f"https://{API_ENDPOINT}/v1beta1/text:synthesize"

        request_body = {
            "input": {"text": text},
            "voice": {
                "language_code": "en-US",
                "voice_clone": {
                    "voice_cloning_key": voice_key,
                },
            },
            "audioConfig": {"audioEncoding": "LINEAR16", "sample_rate_hertz": 24000},
        }

        headers = {
            "Authorization": f"Bearer {ACCESS_TOKEN}",
            "x-goog-user-project": PROJECT_ID,
            "Content-Type": "application/json; charset=utf-8",
        }

        progress(0.6, desc="Processing audio...")
        response = requests.post(url, headers=headers, json=request_body)
        response.raise_for_status()

        response_json = response.json()
        audio_content = response_json.get("audioContent")

        if audio_content:
            progress(1.0, desc="Speech generated!")
            audio_bytes = base64.b64decode(audio_content)
            audio_array = np.frombuffer(audio_bytes, dtype=np.int16)
            return (24000, audio_array), "Speech generated successfully!"
        return None, "Error: Audio content not found in the response."

    except Exception as e:
        return None, f"Error generating speech: {e!s}"


def reset_interface() -> tuple:
    return None, None, "", "", None, None, "Interface reset."

# fmt: off
reference_audio_path = "[your-reference-audio-path]"  # @param {type: "string", placeholder: "[your-reference-audio-path]", isTemplate: true}
consent_audio_path = "[your-consent-audio-path]"  # @param {type: "string", placeholder: "[your-consent-audio-path]", isTemplate: true}
# fmt: on

reference_audio_bytes = wav_to_base64(reference_audio_path)
consent_audio_bytes = wav_to_base64(consent_audio_path)

voice_key = create_instant_custom_voice_key(reference_audio_bytes, consent_audio_bytes)

text_to_synthesize = """
Breaking news! Chirp 3, Google Cloud's audio model, now has Instant Custom Voice.

With Instant Custom Voice, you can generate custom voices with just 10 seconds of audio to empower your AI narration.
Chirp 3 Instant Custom Voice is now available in preview with allowlist.  Check out the link below.

And yes, this voice is generated using Chirp 3 Instant Custom Voice!
"""
synthesize_text_with_cloned_voice(voice_key, text_to_synthesize)

with gr.Blocks(
    theme=gr.themes.Default(
        primary_hue="blue", secondary_hue="blue", neutral_hue="gray"
    )
) as app:
    # Create a state variable to store the full voice key
    full_voice_key = gr.State("")

    # Define title
    gr.Markdown(
        """
        # Chirp 3 - Instant custom voice demo
        Upload reference and consent audio files to create a custom voice, then synthesize speech.
        """
    )

    # Define input and output components
    with gr.Row():
        with gr.Column(scale=1):
            reference_audio = gr.Audio(
                label="Reference Voice", type="filepath", elem_id="reference_audio"
            )
            consent_audio = gr.Audio(
                label="Consent Audio", type="filepath", elem_id="consent_audio"
            )

            create_btn = gr.Button("Create Voice", variant="primary")
            voice_key_output = gr.Textbox(label="Voice Key", elem_id="voice_key")

        with gr.Column(scale=1):
            text_input = gr.Textbox(
                label="Text to Synthesize",
                placeholder="Enter the text you want the voice to say...",
                lines=5,
                elem_id="text_input",
            )
            generate_btn = gr.Button("Generate Speech", variant="primary")
            audio_output = gr.Audio(label="Generated Audio", elem_id="audio_output")
            status_output = gr.Textbox(label="Status", elem_id="status_output")

    with gr.Row():
        clear_btn = gr.Button("Clear All", variant="secondary")

    # Set up event handlers
    create_btn.click(
        create_voice_with_masking,
        inputs=[reference_audio, consent_audio],
        outputs=[full_voice_key, voice_key_output],
    )

    generate_btn.click(
        generate_speech,
        inputs=[full_voice_key, text_input],
        outputs=[audio_output, status_output],
    )

    clear_btn.click(
        reset_interface,
        inputs=[],
        outputs=[
            reference_audio,
            consent_audio,
            voice_key_output,
            full_voice_key,
            text_input,
            audio_output,
            status_output,
        ],
    )

    # Apply custom CSS for Google styling
    gr.Markdown(
        """
        <style>
        .gradio-container {
            font-family: 'Google Sans', 'Roboto', sans-serif !important;
        }
        .gr-button-primary {
            background-color: #4285F4 !important;
        }
        .gr-button-secondary {
            border-color: #4285F4 !important;
            color: #4285F4 !important;
        }
        h1 {
            font-family: 'Google Sans', 'Roboto', sans-serif !important;
            color: #202124 !important;
        }
        </style>
        """
    )

app.launch(share=True)

app.close()