app/ui.py

"""
Gradio frontend for multilingual voice cloning.

Two tabs:
  1. Clone Voice — single language
  2. Language Tour — same reference, many languages

The language dropdown is populated from `VoiceCloneClient.discover_languages()`,
which queries the backend `/languages` endpoint and filters LANGUAGE_CATALOG.
The language code is passed end-to-end into synthesis so the selection
actually changes the output — fixing the bug shown in the original PDF notes.
"""
from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any

import gradio as gr

from app.client import VoiceCloneClient
from app.config import LANGUAGE_CATALOG, get_settings
from app.errors import VoiceCloneError

log = logging.getLogger(__name__)


# Pre-baked demo sentences per language.
DEMO_SENTENCES = {
    "yue": "大家好，歡迎體驗我哋嘅多語言聲音克隆系統，廣東話都支援到。",
    "en":  "Hello everyone, welcome to our multilingual voice cloning demonstration.",
    "zh":  "大家好，欢迎体验我们的多语言语音克隆系统。",
    "ja":  "皆さん、こんにちは。多言語音声クローニングのデモへようこそ。",
    "ko":  "안녕하세요, 다국어 음성 복제 데모에 오신 것을 환영합니다.",
    "fr":  "Bonjour à tous, bienvenue dans notre démo de clonage vocal multilingue.",
    "es":  "Hola a todos, bienvenidos a nuestra demostración de clonación de voz multilingüe.",
}


def _format_label(code: str, meta: dict[str, str]) -> str:
    # Make Cantonese visually prominent.
    if code == "yue":
        return f"🇭🇰  {meta['native']}  — {meta['english']} (yue)"
    return f"{meta['native']} — {meta['english']} ({code})"


def _build_label_map(client: VoiceCloneClient) -> tuple[list[str], dict[str, str]]:
    """Returns (dropdown_choices, {label: code})."""
    langs = client.discover_languages()
    if not langs:
        # Fall back to full catalog so the UI still loads if the backend is cold.
        langs = dict(LANGUAGE_CATALOG)
    # Keep Cantonese pinned to the top regardless of dict order.
    ordered = []
    if "yue" in langs:
        ordered.append(("yue", langs["yue"]))
    for code, meta in langs.items():
        if code != "yue":
            ordered.append((code, meta))
    choices = [_format_label(c, m) for c, m in ordered]
    label_to_code = {_format_label(c, m): c for c, m in ordered}
    return choices, label_to_code


def _format_health(health_map: dict[str, Any]) -> str:
    lines = []
    for name, st in health_map.items():
        mark = "✅" if st.reachable else "⚠️"
        lines.append(f"{mark} **{name}** — {st.detail}")
    return "\n".join(lines)


def build_app(client: VoiceCloneClient | None = None) -> gr.Blocks:
    client = client or VoiceCloneClient()
    settings = client.settings
    choices, label_to_code = _build_label_map(client)
    cantonese_present = "yue" in client.discover_languages()
    default_label = next((lbl for lbl in choices if "yue" in lbl), choices[0] if choices else "")

    # ---- callbacks -----------------------------------------------------
    def on_clone(ref_audio, text, language_label, exaggeration, cfg_weight):
        if not ref_audio:
            return None, "❌ Please upload reference audio first."
        if not language_label:
            return None, "❌ Please select a language."
        code = label_to_code.get(language_label)
        if not code:
            return None, f"❌ Unknown language: {language_label}"
        try:
            res = client.clone_voice(
                reference_audio_path=Path(ref_audio),
                text=text,
                language_code=code,
                exaggeration=float(exaggeration),
                cfg_weight=float(cfg_weight),
            )
        except VoiceCloneError as e:
            return None, f"❌ {e}"
        status = (
            f"✅ Generated in {res.latency_seconds:.1f}s via **{res.provider_name}** "
            f"— language={code} ({language_label}), model={res.model}, "
            f"sr={res.sample_rate} Hz"
        )
        return str(res.wav_path), status

    def on_tour(ref_audio, text, exaggeration, cfg_weight):
        if not ref_audio:
            return [None] * len(choices) + ["❌ Please upload reference audio first."]
        results = client.language_tour(
            reference_audio_path=Path(ref_audio),
            text=text,
            exaggeration=float(exaggeration),
            cfg_weight=float(cfg_weight),
            language_codes=[label_to_code[lbl] for lbl in choices],
        )
        audio_outputs: list[str | None] = []
        lines = []
        for (code, res, err) in results:
            meta = LANGUAGE_CATALOG.get(code, {"english": code, "native": code})
            if res:
                audio_outputs.append(str(res.wav_path))
                lines.append(f"✅ {meta['english']} ({code}) — {res.latency_seconds:.1f}s via {res.provider_name}")
            else:
                audio_outputs.append(None)
                lines.append(f"❌ {meta['english']} ({code}) — {err}")
        return audio_outputs + ["\n".join(lines)]

    def on_health():
        return _format_health(client.health())

    def on_language_change(label: str) -> str:
        code = label_to_code.get(label or "", "en")
        return DEMO_SENTENCES.get(code, "")

    # ---- layout --------------------------------------------------------
    with gr.Blocks(
        title="ZhouFenMianFan — Multilingual Voice Cloning (Cloud)",
        theme=gr.themes.Soft(),
    ) as demo:
        gr.Markdown(
            "# 🎤 粥粉麵飯 — Multilingual Voice Cloning\n"
            "Cross-lingual voice cloning for Hong Kong users. Upload a "
            "**3–30 second reference clip** and generate speech in 粵語 "
            "(Cantonese), English, Mandarin, Japanese, Korean, French, "
            "and Spanish. Powered by a Chatterbox-compatible cloud API "
            "with a MiniMax adapter for Cantonese.\n"
        )
        if not cantonese_present:
            gr.Markdown(
                "> ⚠️ **Cantonese (yue) is currently unavailable.** "
                "Configure `MINIMAX_API_KEY` in your `.env` to enable it."
            )

        with gr.Accordion("Backend status", open=False):
            health_md = gr.Markdown(_format_health(client.health()))
            refresh_btn = gr.Button("Refresh health", size="sm")
            refresh_btn.click(on_health, outputs=health_md)

        with gr.Tab("Clone Voice"):
            with gr.Row():
                with gr.Column():
                    ref1 = gr.Audio(
                        type="filepath",
                        label=f"Reference voice (WAV, "
                              f"{settings.min_reference_seconds:.0f}–"
                              f"{settings.max_reference_seconds:.0f}s)",
                        sources=["upload", "microphone"],
                    )
                    lang1 = gr.Dropdown(
                        choices=choices,
                        value=default_label,
                        label="Target language",
                        info="Cantonese (粵語) is routed through the MiniMax adapter.",
                    )
                    text1 = gr.Textbox(
                        lines=3,
                        value=DEMO_SENTENCES.get("yue", "") if cantonese_present
                              else DEMO_SENTENCES.get("en", ""),
                        label="Text to synthesise",
                    )
                    with gr.Row():
                        exag1 = gr.Slider(
                            settings.exaggeration_range[0],
                            settings.exaggeration_range[1],
                            value=settings.default_exaggeration, step=0.05,
                            label="Exaggeration",
                        )
                        cfg1 = gr.Slider(
                            settings.cfg_weight_range[0],
                            settings.cfg_weight_range[1],
                            value=settings.default_cfg_weight, step=0.05,
                            label="CFG weight",
                        )
                    go1 = gr.Button("🎙️ Generate", variant="primary")
                with gr.Column():
                    audio1 = gr.Audio(type="filepath", label="Cloned output",
                                      interactive=False)
                    status1 = gr.Markdown("")

            lang1.change(on_language_change, inputs=lang1, outputs=text1)
            go1.click(
                on_clone,
                inputs=[ref1, text1, lang1, exag1, cfg1],
                outputs=[audio1, status1],
            )

        with gr.Tab("Language Tour"):
            gr.Markdown(
                "Clone the same reference voice across every supported language "
                "using one shared sentence."
            )
            with gr.Row():
                ref2 = gr.Audio(type="filepath",
                                label="Reference voice (WAV, 3–30s)",
                                sources=["upload", "microphone"])
                text2 = gr.Textbox(
                    lines=2,
                    value="Hello, welcome to our AI voice cloning demonstration.",
                    label="Shared text (model will synthesise this same "
                          "sentence in every language)",
                )
            with gr.Row():
                exag2 = gr.Slider(
                    settings.exaggeration_range[0],
                    settings.exaggeration_range[1],
                    value=settings.default_exaggeration, step=0.05,
                    label="Exaggeration",
                )
                cfg2 = gr.Slider(
                    settings.cfg_weight_range[0],
                    settings.cfg_weight_range[1],
                    value=settings.default_cfg_weight, step=0.05,
                    label="CFG weight",
                )
            go2 = gr.Button("🌏 Generate language tour", variant="primary")
            audio_widgets: list[gr.Audio] = []
            for lbl in choices:
                audio_widgets.append(gr.Audio(type="filepath", label=lbl,
                                              interactive=False))
            status2 = gr.Markdown("")

            go2.click(
                on_tour,
                inputs=[ref2, text2, exag2, cfg2],
                outputs=audio_widgets + [status2],
            )

        gr.Markdown(
            "---\n"
            "**Guidance**: 10–30 s of clean speech works best. Avoid music, "
            "background noise, and clipping. 24 kHz / 32 kHz mono WAV "
            "produces the cleanest clones. Long-form texts over 500 characters "
            "are split by the backend automatically.\n"
        )
    return demo


def main() -> None:
    logging.basicConfig(
        level=os.getenv("LOG_LEVEL", "INFO"),
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    settings = get_settings()
    demo = build_app()
    demo.queue().launch(
        server_name=settings.gradio_server_name,
        server_port=settings.gradio_server_port,
        share=False,
        show_error=True,
    )


if __name__ == "__main__":
    main()