from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import tempfile
import uuid
from pathlib import Path

import gradio as gr
from pyharp import ModelCard, build_endpoint
from pyharp.tags import Channels, SampleRate, Subcategory

from pocket_runtime import load_model, synthesize

MAX_TEXT_LENGTH = 500
OUTPUT_ROOT = Path(tempfile.gettempdir()) / "pocket_tts_outputs"

model_card = ModelCard(
    name="Pocket TTS",
    description=(
        "Generate English speech with Kyutai's compact Pocket TTS model using "
        "CPU inference and predefined voices, as 24 kHz mono audio. The code is MIT, but the "
        "model weights are CC-BY-4.0, and their model card prohibits uses such as cloning or "
        "impersonating a voice without consent, deception, and harmful content. Repository: "
        "https://github.com/kyutai-labs/pocket-tts Paper: "
        "https://arxiv.org/abs/2509.06926"
    ),
    author="Simon Rouard, Manu Orsini, Axel Roebel, et al.",
    tags=[Subcategory.TEXT_TO_SPEECH, SampleRate(24000), Channels(1), 'streaming-tts', 'domain:speech'],
)


def process_fn(text: str | None, voice: str) -> str:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    clean_text = (text or "").strip()
    if not clean_text:
        raise gr.Error("Please enter text to synthesize.")
    if len(clean_text) > MAX_TEXT_LENGTH:
        raise gr.Error(
            f"Text must be no longer than {MAX_TEXT_LENGTH} characters. "
            f"Received {len(clean_text)} characters."
        )

    return speak(clean_text, voice)


@spaces.GPU
def speak(clean_text: str, voice: str) -> str:
    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "pocket_tts_output.wav"
    try:
        synthesize(clean_text, voice, output_path)
    except Exception as exc:
        raise gr.Error(f"Pocket TTS synthesis failed: {exc}") from exc
    return str(output_path)


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is loaded there, before the first request rather than during it.
if __name__ != "__main__":
    load_model()


if __name__ == "__main__":
    with gr.Blocks(title="Pocket TTS") as demo:
        input_components = [
            gr.Textbox(
                label="English Text",
                lines=4,
                max_lines=8,
                placeholder="Enter the sentence you want to hear...",
            )
            .harp_required(True)
            .set_info("English text, up to 500 characters."),
            gr.Dropdown(
                choices=[
                    "alba",
                    "anna",
                    "bill_boerst",
                    "eve",
                    "george",
                    "jane",
                    "mary",
                    "michael",
                    "paul",
                    "stuart_bell",
                    "vera",
                ],
                value="alba",
                label="Voice",
            ).set_info("A predefined English voice supplied by the model authors."),
        ]
        output_components = [
            gr.Audio(
                type="filepath",
                label="Synthesized Speech",
            ).set_info("24 kHz mono WAV generated on CPU."),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )


    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)