#!/usr/bin/env python3
#
# Copyright      2022-2023  Xiaomi Corp.        (authors: Fangjun Kuang)
#
# See LICENSE for clarification regarding multiple authors
#
# 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
#
#     http://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.

# References:
# https://gradio.app/docs/#dropdown

import os
import tempfile
import time
import urllib.request
import uuid
from datetime import datetime
from pathlib import Path

import gradio as gr
import librosa
import numpy as np
import sherpa_onnx
import soundfile as sf

from model import get_pretrained_model, language_to_models, language_to_supertonic_lang

hint_filename = f"/tmp/{uuid.uuid4()}.wav"
os.system(
    f"""
wget -O {hint_filename} https://huggingface.co/spaces/k2-fsa/text-to-speech/resolve/main/hint.wav
ls -lh *.wav
ls -lh /tmp/*.wav
"""
)


def MyPrint(s):
    now = datetime.now()
    date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
    print(f"{date_time}: {s}")


def convert_to_wav(in_filename: str) -> str:
    """Convert the input audio file to a wave file using ffmpeg"""
    out_filename = f"{in_filename}.wav"
    MyPrint(f"Converting '{in_filename}' to '{out_filename}'")
    os.system(
        f"ffmpeg -hide_banner -loglevel error -i '{in_filename}' -ar 16000 -ac 1 '{out_filename}' -y"
    )
    return out_filename


title = "# Next-gen Kaldi: Text-to-speech (TTS)"

description = """
This space shows how to convert text to speech with Next-gen Kaldi.

It is running on CPU within a docker container provided by Hugging Face.

**Voice Cloning**: Select "Voice Cloning" language to use voice cloning models:
- [Pocket TTS](https://github.com/kyutai-labs/pocket-tts): Supports **6 languages** (English, French, German, Portuguese, Italian, Spanish). Only requires a reference audio clip.
- [ZipVoice](https://github.com/k2-fsa/zipvoice): Supports **Chinese and English**. Requires both a reference audio clip and the exact text spoken in the reference audio.

You need to provide a reference audio clip (upload, record, or URL) to clone the voice.

See more information by visiting the following links:

- <https://github.com/k2-fsa/sherpa-onnx>
- <https://github.com/kyutai-labs/pocket-tts>
- <https://github.com/k2-fsa/zipvoice>
- <https://k2-fsa.github.io/sherpa/onnx/tts/pocket.html>
- <https://k2-fsa.github.io/sherpa/onnx/tts/zipvoice.html>

If you want to deploy it locally, please see
<https://k2-fsa.github.io/sherpa/>

If you want to use Android APKs, please see
<https://k2-fsa.github.io/sherpa/onnx/tts/apk.html>

If you want to use Android text-to-speech engine APKs, please see
<https://k2-fsa.github.io/sherpa/onnx/tts/apk-engine.html>

If you want to download an all-in-one exe for Windows, please see
<https://github.com/k2-fsa/sherpa-onnx/releases/tag/tts-models>

See also <https://k2-fsa.github.io/sherpa/onnx/tts/all/>
for models with audio samples.

"""

# css style is copied from
# https://huggingface.co/spaces/alphacep/asr/blob/main/app.py#L113
css = """
.result {display:flex;flex-direction:column}
.result_item {padding:15px;margin-bottom:8px;border-radius:15px;width:100%}
.result_item_success {background-color:mediumaquamarine;color:white;align-self:start}
.result_item_error {background-color:#ff7070;color:white;align-self:start}
"""


def build_html_output(s: str, style: str = "result_item_success"):
    return f"""
    <div class='result'>
        <div class='result_item {style}'>
          {s}
        </div>
    </div>
    """


def process(
    language: str,
    repo_id: str,
    text: str,
    sid: str,
    speed: float,
    reference_audio_path: str = None,
    reference_text: str = None,
):
    max_len = 4000
    MyPrint(f"Input text {len(text)}: {text[:max_len]}. sid: {sid}, speed: {speed}")
    if len(text) > max_len:
        MyPrint(f"Too long! {len(text)}")
        info = """
            To ensure this space is responsive, please use short text for testing.<br/>
            You can run this space locally to process long texts.<br>
            See also https://k2-fsa.github.io/sherpa/onnx/
        """
        return hint_filename, build_html_output(info)

    if sid is None or str(sid).strip() == "":
        sid = 0
    else:
        sid = int(sid)
    tts = get_pretrained_model(repo_id, speed)

    start = time.time()
    if "pocket" in repo_id:
        if reference_audio_path is None or not Path(reference_audio_path).is_file():
            info = "Please provide a reference audio for voice cloning."
            return hint_filename, build_html_output(info, "result_item_error")

        # Convert reference audio to wav format
        ref_wav = convert_to_wav(reference_audio_path)
        MyPrint(f"Loading reference audio: {ref_wav}")

        # Check reference audio duration (max 60 seconds)
        ref_info = sf.info(ref_wav)
        ref_duration = ref_info.duration
        MyPrint(f"Reference audio duration: {ref_duration:.2f} seconds")
        if ref_duration > 60:
            info = f"""
                Reference audio is too long ({ref_duration:.1f} seconds).<br/>
                We accept only reference audio up to 60 seconds.<br/>
                Please provide a shorter reference audio.
            """
            return hint_filename, build_html_output(info, "result_item_error")

        reference_audio, sample_rate = librosa.load(ref_wav, sr=tts.sample_rate)

        gen_config = sherpa_onnx.GenerationConfig()
        gen_config.reference_audio = reference_audio
        gen_config.reference_sample_rate = sample_rate
        gen_config.num_steps = 5
        audio = tts.generate(text, gen_config)
    elif "zipvoice" in repo_id:
        if reference_audio_path is None or not Path(reference_audio_path).is_file():
            info = "Please provide a reference audio for voice cloning."
            return hint_filename, build_html_output(info, "result_item_error")

        if reference_text is None or reference_text.strip() == "":
            info = "Please provide the text content of the reference audio."
            return hint_filename, build_html_output(info, "result_item_error")

        # Convert reference audio to wav format
        ref_wav = convert_to_wav(reference_audio_path)
        MyPrint(f"Loading reference audio: {ref_wav}")

        # Check reference audio duration (max 60 seconds)
        ref_info = sf.info(ref_wav)
        ref_duration = ref_info.duration
        MyPrint(f"Reference audio duration: {ref_duration:.2f} seconds")
        if ref_duration > 60:
            info = f"""
                Reference audio is too long ({ref_duration:.1f} seconds).<br/>
                We accept only reference audio up to 60 seconds.<br/>
                Please provide a shorter reference audio.
            """
            return hint_filename, build_html_output(info, "result_item_error")

        reference_audio, sample_rate = librosa.load(ref_wav, sr=None)

        gen_config = sherpa_onnx.GenerationConfig()
        gen_config.reference_audio = reference_audio
        gen_config.reference_sample_rate = sample_rate
        gen_config.reference_text = reference_text
        gen_config.num_steps = 4
        gen_config.extra["min_char_in_sentence"] = "30"
        audio = tts.generate(text, gen_config)
    elif "supertonic" in repo_id:
        gen_config = sherpa_onnx.GenerationConfig()
        gen_config.sid = sid
        gen_config.num_steps = 8
        gen_config.speed = speed
        lang_code = language_to_supertonic_lang.get(language, "en")
        gen_config.extra["lang"] = lang_code
        audio = tts.generate(text, gen_config)
    else:
        gen_config = sherpa_onnx.GenerationConfig()
        gen_config.sid = sid
        gen_config.speed = speed
        audio = tts.generate(text, gen_config)
    end = time.time()

    if len(audio.samples) == 0:
        raise ValueError(
            "Error in generating audios. Please read previous error messages."
        )

    duration = len(audio.samples) / audio.sample_rate

    elapsed_seconds = end - start
    rtf = elapsed_seconds / duration

    info = f"""
    Wave duration  : {duration:.3f} s <br/>
    Processing time: {elapsed_seconds:.3f} s <br/>
    RTF: {elapsed_seconds:.3f}/{duration:.3f} = {rtf:.3f} <br/>
    """

    MyPrint(info)
    MyPrint(f"\nrepo_id: {repo_id}\ntext: {text}\nsid: {sid}\nspeed: {speed}")

    samples = np.array(audio.samples)
    if samples.dtype == np.float64 or samples.dtype == np.float32:
        samples = (samples * 32767).astype(np.int16)
    return (audio.sample_rate, samples), build_html_output(info)


def update_model_dropdown(language: str):
    if language in language_to_models:
        choices = language_to_models[language]
        return gr.Dropdown(
            choices=choices,
            value=choices[0],
            interactive=True,
        )

    raise ValueError(f"Unsupported language: {language}")


def toggle_visibility(language: str):
    """Show reference audio section only for Voice Cloning"""
    is_voice_cloning = language == "Voice Cloning"
    return (
        gr.update(visible=is_voice_cloning),  # ref_audio_section
        gr.update(visible=not is_voice_cloning),  # input_sid
        gr.update(visible=not is_voice_cloning),  # input_button
    )


def toggle_ref_text_visibility(repo_id: str):
    """Show reference text only for ZipVoice model"""
    is_zipvoice = "zipvoice" in repo_id.lower() if repo_id else False
    return gr.update(visible=is_zipvoice)


def check_ref_submit_state(
    language: str,
    repo_id: str,
    text: str,
    ref_upload: str,
    ref_microphone: str,
    ref_url: str,
    ref_text: str,
):
    """Enable submit button only when all required inputs are provided"""
    has_text = text is not None and text.strip() != ""

    # Determine which reference audio source is active
    ref_audio = (
        ref_upload
        or ref_microphone
        or (ref_url if ref_url and ref_url.strip() else None)
    )
    has_ref_audio = ref_audio is not None

    is_zipvoice = "zipvoice" in repo_id.lower() if repo_id else False
    is_pocket = "pocket" in repo_id.lower() if repo_id else False

    if is_zipvoice:
        # ZipVoice requires reference audio, reference text, and input text
        has_ref_text = ref_text is not None and ref_text.strip() != ""
        enabled = has_text and has_ref_audio and has_ref_text
    elif is_pocket:
        # Pocket TTS requires reference audio and input text
        enabled = has_text and has_ref_audio
    else:
        # Other models only need input text
        enabled = has_text

    return gr.update(interactive=enabled)


demo = gr.Blocks(css=css)


with demo:
    gr.Markdown(title)
    language_choices = list(language_to_models.keys())

    language_radio = gr.Radio(
        label="Language",
        choices=language_choices,
        value=language_choices[0],
    )

    model_dropdown = gr.Dropdown(
        choices=language_to_models[language_choices[0]],
        label="Select a model",
        value=language_to_models[language_choices[0]][0],
    )

    language_radio.change(
        update_model_dropdown,
        inputs=language_radio,
        outputs=model_dropdown,
    )

    with gr.Tabs():
        with gr.TabItem("Please input your text"):
            input_text = gr.Textbox(
                label="Input text",
                info="Your text",
                lines=3,
                placeholder="Please input your text here",
            )

            input_sid = gr.Textbox(
                label="Speaker ID",
                info="Speaker ID",
                lines=1,
                max_lines=1,
                value="0",
                placeholder="Speaker ID. Valid only for multi-speaker model",
                visible=False,  # Hidden by default since Voice Cloning is first
            )

            input_speed = gr.Slider(
                minimum=0.1,
                maximum=10,
                value=1,
                step=0.1,
                label="Speed (larger->faster; smaller->slower)",
            )

            # Voice cloning reference audio section (visible by default since Voice Cloning is first)
            with gr.Column(visible=True) as ref_audio_section:
                gr.Markdown("### Reference Audio for Voice Cloning")
                gr.Markdown(
                    "Provide a reference audio clip. The generated speech will "
                    "clone the voice from this audio."
                )

                with gr.Tabs():
                    with gr.TabItem("Upload from disk"):
                        ref_upload = gr.Audio(
                            sources=["upload"],
                            type="filepath",
                            label="Upload reference audio",
                        )

                    with gr.TabItem("Record from microphone"):
                        ref_microphone = gr.Audio(
                            sources=["microphone"],
                            type="filepath",
                            label="Record reference audio",
                        )

                    with gr.TabItem("From URL"):
                        ref_url = gr.Textbox(
                            max_lines=1,
                            placeholder="URL to a reference audio file",
                            label="Reference audio URL",
                            interactive=True,
                        )
                        ref_url_audio = gr.Audio(label="Downloaded reference audio")

                # Reference text section (for ZipVoice, hidden by default)
                ref_text_box = gr.Textbox(
                    label="Reference Text",
                    info="The text content of the reference audio (required for ZipVoice). "
                    "Must match exactly what is said in the reference audio.",
                    lines=2,
                    placeholder="Enter the exact text spoken in the reference audio...",
                    visible=False,  # Hidden by default, shown only for ZipVoice
                )

                # Submit button for voice cloning
                ref_submit_button = gr.Button("Submit", interactive=False)

            # Normal submit button (hidden for voice cloning)
            input_button = gr.Button(
                "Submit", visible=False
            )  # Hidden by default since Voice Cloning is first

            output_audio = gr.Audio(label="Output", type="numpy")

            output_info = gr.HTML(label="Info")

    # Toggle visibility based on language selection
    language_radio.change(
        toggle_visibility,
        inputs=language_radio,
        outputs=[ref_audio_section, input_sid, input_button],
    )

    # Toggle reference text visibility based on model selection
    model_dropdown.change(
        toggle_ref_text_visibility,
        inputs=model_dropdown,
        outputs=ref_text_box,
    )

    # Check state and enable/disable voice cloning submit button
    # We need to check when any input changes
    ref_inputs = [
        language_radio,
        model_dropdown,
        input_text,
        ref_upload,
        ref_microphone,
        ref_url,
        ref_text_box,
    ]

    for inp in ref_inputs:
        inp.change(
            check_ref_submit_state,
            inputs=ref_inputs,
            outputs=ref_submit_button,
        )

    # Normal submit button (non-voice-cloning models)
    input_button.click(
        process,
        inputs=[
            language_radio,
            model_dropdown,
            input_text,
            input_sid,
            input_speed,
        ],
        outputs=[
            output_audio,
            output_info,
        ],
    )

    # Voice cloning submit button
    def process_voice_clone(
        language,
        repo_id,
        text,
        sid,
        speed,
        ref_upload,
        ref_microphone,
        ref_url,
        ref_text,
    ):
        """Process voice cloning based on which audio source is provided"""
        ref_audio = ref_upload or ref_microphone
        if ref_audio is None and ref_url and ref_url.strip():
            # Download from URL
            MyPrint(f"Downloading reference audio from URL: {ref_url}")
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                try:
                    urllib.request.urlretrieve(ref_url, f.name)
                    ref_audio = f.name
                except Exception as e:
                    MyPrint(f"Error downloading URL: {e}")
                    return (
                        hint_filename,
                        build_html_output(
                            f"Error downloading audio: {e}", "result_item_error"
                        ),
                        None,
                    )

        if ref_audio is None:
            return (
                hint_filename,
                build_html_output(
                    "Please provide a reference audio.", "result_item_error"
                ),
                None,
            )

        output_audio, output_info = process(
            language,
            repo_id,
            text,
            sid,
            speed,
            reference_audio_path=ref_audio,
            reference_text=ref_text,
        )
        return output_audio, output_info, ref_audio if ref_url else None

    ref_submit_button.click(
        process_voice_clone,
        inputs=[
            language_radio,
            model_dropdown,
            input_text,
            input_sid,
            input_speed,
            ref_upload,
            ref_microphone,
            ref_url,
            ref_text_box,
        ],
        outputs=[
            output_audio,
            output_info,
            ref_url_audio,
        ],
    )

    gr.Markdown(description)


def download_espeak_ng_data():
    os.system(
        """
    cd /tmp
    wget -qq https://github.com/k2-fsa/sherpa-onnx/releases/download/tts-models/espeak-ng-data.tar.bz2
    tar xf espeak-ng-data.tar.bz2
    """
    )
    if not Path("/tmp/dict").is_dir():
        os.system(
            "cd /tmp; curl -SL -O https://github.com/csukuangfj/cppjieba/releases/download/sherpa-onnx-2024-04-19/dict.tar.bz2; tar xvf dict.tar.bz2"
        )
    os.system("ls -lh /tmp/dict")


if __name__ == "__main__":
    download_espeak_ng_data()
    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"

    demo.launch()