import gradio as gr
import os
import shutil
import tempfile
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import FileResponse

# ── Singletons séparés ────────────────────────────────────────────────────────
_kokoro = None   # léger, chargé au premier /api/speak ou Simple TTS
_cloner = None   # lourd (840MB), chargé seulement si clone/convert demandé

def get_kokoro():
    """Charge uniquement Kokoro ONNX — sans Kanade."""
    global _kokoro
    if _kokoro is None:
        from core.cloner import KokoClone
        print("[KokoOnly] Loading Kokoro TTS (no Kanade)...")
        _kokoro = KokoClone.__new__(KokoClone)
        # Init manuelle sans Kanade
        import torch
        _kokoro.device = torch.device("cpu")
        _kokoro.kokoro_cache = {}
        _kokoro.hf_repo = "PatnaikAshish/kokoclone"
        print("[KokoOnly] Ready.")
    return _kokoro

def get_cloner():
    """Charge KokoClone complet avec Kanade — seulement si nécessaire."""
    global _cloner
    if _cloner is None:
        from core.cloner import KokoClone
        print("[KokoClone] Loading full model (Kokoro + Kanade)...")
        _cloner = KokoClone()
        print("[KokoClone] Ready.")
    return _cloner


# ── Handlers Gradio ───────────────────────────────────────────────────────────

def gradio_speak(text, lang, speed):
    if not text or not text.strip():
        raise gr.Error("Veuillez entrer du texte.")
    output_file = "/tmp/gradio_speak_output.wav"
    try:
        get_cloner().speak(text=text, lang=lang, output_path=output_file, speed=float(speed))
        return output_file
    except Exception as e:
        raise gr.Error(f"Erreur : {str(e)}")


def gradio_clone(text, lang, ref_audio_path):
    if not text or not text.strip():
        raise gr.Error("Please enter some text.")
    if not ref_audio_path:
        raise gr.Error("Please upload or record a reference audio file.")
    output_file = "/tmp/gradio_clone_output.wav"
    try:
        get_cloner().generate(text=text, lang=lang, reference_audio=ref_audio_path, output_path=output_file)
        return output_file
    except Exception as e:
        raise gr.Error(f"An error occurred: {str(e)}")


def gradio_convert(source_audio_path, ref_audio_path):
    if not source_audio_path:
        raise gr.Error("Please upload a source audio file.")
    if not ref_audio_path:
        raise gr.Error("Please upload a reference audio file.")
    output_file = "/tmp/gradio_convert_output.wav"
    try:
        get_cloner().convert(source_audio=source_audio_path, reference_audio=ref_audio_path, output_path=output_file)
        return output_file
    except Exception as e:
        raise gr.Error(f"An error occurred: {str(e)}")


# ── Interface Gradio ──────────────────────────────────────────────────────────

with gr.Blocks() as demo:
    gr.Markdown(
        """
        <div style="text-align: center;">
            <h1>🎧 KokoClone</h1>
            <p>Synthèse vocale multilingue · Voice Cloning · Re-voicing</p>
        </div>
        """
    )

    with gr.Tabs():

        with gr.Tab("🗣️ Simple TTS"):
            gr.Markdown("Génère la parole avec une voix Kokoro — **aucun audio de référence requis**.")
            with gr.Row():
                with gr.Column(scale=1):
                    speak_text = gr.Textbox(
                        label="Texte à lire",
                        lines=5,
                        placeholder="Hello! How are you today?\nBonjour ! Comment allez-vous ?"
                    )
                    speak_lang = gr.Dropdown(
                        label="Langue / Voix",
                        choices=[
                            ("🇬🇧 English — af_bella", "en"),
                            ("🇫🇷 Français — ff_siwis", "fr"),
                            ("🇮🇳 Hindi — hf_alpha", "hi"),
                            ("🇪🇸 Español — im_nicola", "es"),
                            ("🇮🇹 Italiano — im_nicola", "it"),
                            ("🇧🇷 Português — pf_dora", "pt"),
                        ],
                        value="en"
                    )
                    speak_speed = gr.Slider(
                        label="Vitesse",
                        minimum=0.5,
                        maximum=1.5,
                        value=1.0,
                        step=0.1
                    )
                    speak_btn = gr.Button("🗣️ Générer l'audio", variant="primary")
                with gr.Column(scale=1):
                    speak_out = gr.Audio(label="Audio généré", interactive=False, autoplay=True)
                    gr.Markdown(
                        """
**Voix disponibles**
- 🇬🇧 `af_bella` — femme américaine, claire
- 🇫🇷 `ff_siwis` — femme française, fluide
- 🇮🇳 `hf_alpha` — Hindi
- 🇪🇸🇮🇹 `im_nicola` — voix masculine
- 🇧🇷 `pf_dora` — Portugais brésilien

> Vitesse **0.7–0.8** recommandée pour les débutants.
                        """
                    )
            speak_btn.click(
                fn=gradio_speak,
                inputs=[speak_text, speak_lang, speak_speed],
                outputs=speak_out
            )

        with gr.Tab("🎤 Text → Clone"):
            gr.Markdown("Synthétise le texte puis applique la voix d'un audio de référence.")
            with gr.Row():
                with gr.Column(scale=1):
                    clone_text = gr.Textbox(label="1. Texte à synthétiser", lines=4)
                    clone_lang = gr.Dropdown(
                        label="2. Langue",
                        choices=[
                            ("English", "en"), ("Français", "fr"), ("Hindi", "hi"),
                            ("Japanese", "ja"), ("Chinese", "zh"), ("Español", "es"),
                            ("Italiano", "it"), ("Português", "pt"),
                        ],
                        value="en"
                    )
                    clone_ref = gr.Audio(label="3. Audio de référence", type="filepath")
                    clone_btn = gr.Button("🚀 Générer avec clonage", variant="primary")
                with gr.Column(scale=1):
                    clone_out = gr.Audio(label="Audio cloné", interactive=False, autoplay=False)
                    gr.Markdown(
                        """
**Conseils**
- Audio de référence propre, sans bruit de fond
- 3 à 10 secondes = résultat optimal
- La langue doit correspondre au texte
                        """
                    )
            clone_btn.click(
                fn=gradio_clone,
                inputs=[clone_text, clone_lang, clone_ref],
                outputs=clone_out
            )

        with gr.Tab("🔁 Audio → Clone"):
            gr.Markdown("Re-voice un audio existant pour qu'il sonne comme la voix de référence.")
            with gr.Row():
                with gr.Column(scale=1):
                    convert_src = gr.Audio(label="1. Audio source", type="filepath")
                    convert_ref = gr.Audio(label="2. Voix de référence", type="filepath")
                    convert_btn = gr.Button("🔁 Convertir la voix", variant="primary")
                with gr.Column(scale=1):
                    convert_out = gr.Audio(label="Audio converti", interactive=False, autoplay=False)
                    gr.Markdown(
                        """
**Comment ça marche**
- **Source** : n'importe quel enregistrement vocal
- **Référence** : court clip du locuteur cible
- KokoClone re-voice la source vers la référence
- Aucune transcription nécessaire
                        """
                    )
            convert_btn.click(
                fn=gradio_convert,
                inputs=[convert_src, convert_ref],
                outputs=convert_out
            )


# ── FastAPI ───────────────────────────────────────────────────────────────────

app = gr.mount_gradio_app(FastAPI(), demo, path="/")


@app.get("/api/voices")
def list_voices():
    return {
        "voices": [
            {"lang": "en", "voice": "af_bella",  "description": "English - American Female"},
            {"lang": "fr", "voice": "ff_siwis",  "description": "French - Female"},
            {"lang": "hi", "voice": "hf_alpha",  "description": "Hindi - Female"},
            {"lang": "es", "voice": "im_nicola", "description": "Spanish - Male"},
            {"lang": "it", "voice": "im_nicola", "description": "Italian - Male"},
            {"lang": "pt", "voice": "pf_dora",   "description": "Portuguese BR - Female"},
        ]
    }


@app.post("/api/speak")
async def api_speak(
    text: str = Form(...),
    lang: str = Form(default="en"),
    speed: float = Form(default=1.0),
):
    out_path = tempfile.mktemp(suffix=".wav")
    try:
        get_cloner().speak(text=text, lang=lang, output_path=out_path, speed=speed)
        return FileResponse(out_path, media_type="audio/wav", filename="speech.wav")
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/clone")
async def api_clone(
    text: str = Form(...),
    lang: str = Form(default="en"),
    reference_audio: UploadFile = File(...),
):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as ref_tmp:
        shutil.copyfileobj(reference_audio.file, ref_tmp)
        ref_path = ref_tmp.name
    out_path = ref_path.replace(".wav", "_out.wav")
    try:
        get_cloner().generate(text=text, lang=lang, reference_audio=ref_path, output_path=out_path)
        return FileResponse(out_path, media_type="audio/wav", filename="cloned.wav")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(ref_path):
            os.unlink(ref_path)


@app.post("/api/convert")
async def api_convert(
    source_audio: UploadFile = File(...),
    reference_audio: UploadFile = File(...),
):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as src_tmp:
        shutil.copyfileobj(source_audio.file, src_tmp)
        src_path = src_tmp.name
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as ref_tmp:
        shutil.copyfileobj(reference_audio.file, ref_tmp)
        ref_path = ref_tmp.name
    out_path = src_path.replace(".wav", "_converted.wav")
    try:
        get_cloner().convert(source_audio=src_path, reference_audio=ref_path, output_path=out_path)
        return FileResponse(out_path, media_type="audio/wav", filename="converted.wav")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(src_path):
            os.unlink(src_path)
        if os.path.exists(ref_path):
            os.unlink(ref_path)