"""
Noukiko - moteur de doublage automatique
Vidéo (anglais/français) -> transcription -> traduction (Fon, Yoruba...) -> voix -> vidéo doublée.

API :
  POST /dub        {"video_url": "...", "source_lang": "en", "target_lang": "fon"}  -> {"job_id": "..."}
  GET  /jobs/{id}  -> {"status": "en_attente|en_cours|termine|erreur", "progress": 0-100,
                       "result_url": "...", "error": "..."}
  GET  /files/{id}.mp4  -> la vidéo doublée
  GET  /languages  -> langues disponibles
"""
import os
import re
import uuid
import shutil
import subprocess
import tempfile
import threading
from concurrent.futures import ThreadPoolExecutor

import requests
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel

# ---------------------------------------------------------------- config
OUT_DIR = os.environ.get("NOUKIKO_OUT", "/tmp/noukiko_out")
PUBLIC_URL = os.environ.get("PUBLIC_URL", "https://hinnou-noukiko-api.hf.space").rstrip("/")  # ex: https://xxx.hf.space
MAX_SECONDS = int(os.environ.get("MAX_SECONDS", "600"))     # durée max d'une vidéo
ORIGINAL_VOLUME = float(os.environ.get("ORIGINAL_VOLUME", "0.12"))  # son d'origine en fond
WHISPER_SIZE = os.environ.get("WHISPER_SIZE", "small")
NLLB_MODEL = os.environ.get("NLLB_MODEL", "facebook/nllb-200-distilled-600M")
os.makedirs(OUT_DIR, exist_ok=True)

# Langues sources (Whisper) -> code NLLB
SOURCE_LANGS = {"en": "eng_Latn", "fr": "fra_Latn"}

# Langues cibles : code NLLB + modèle de voix Meta MMS-TTS
TARGET_LANGS = {
    "fon": {"name": "Fon", "nllb": "fon_Latn", "tts": "facebook/mms-tts-fon"},
    "yor": {"name": "Yoruba", "nllb": "yor_Latn", "tts": "facebook/mms-tts-yor"},
    "ewe": {"name": "Ewe", "nllb": "ewe_Latn", "tts": "facebook/mms-tts-ewe"},
    "hau": {"name": "Haoussa", "nllb": "hau_Latn", "tts": "facebook/mms-tts-hau"},
    "fra": {"name": "Français", "nllb": "fra_Latn", "tts": "facebook/mms-tts-fra"},
    "eng": {"name": "Anglais", "nllb": "eng_Latn", "tts": "facebook/mms-tts-eng"},
}

# ---------------------------------------------------------------- modèles (chargés à la demande)
_lock = threading.Lock()
_models = {}


def get_whisper():
    with _lock:
        if "whisper" not in _models:
            from faster_whisper import WhisperModel
            _models["whisper"] = WhisperModel(WHISPER_SIZE, device="cpu", compute_type="int8")
        return _models["whisper"]


def get_translator():
    with _lock:
        if "nllb" not in _models:
            from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
            tok = AutoTokenizer.from_pretrained(NLLB_MODEL)
            mdl = AutoModelForSeq2SeqLM.from_pretrained(NLLB_MODEL)
            _models["nllb"] = (tok, mdl)
        return _models["nllb"]


def get_tts(model_id):
    with _lock:
        key = "tts:" + model_id
        if key not in _models:
            from transformers import VitsModel, AutoTokenizer
            _models[key] = (AutoTokenizer.from_pretrained(model_id), VitsModel.from_pretrained(model_id))
        return _models[key]


# ---------------------------------------------------------------- étapes
def run(cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        raise RuntimeError(f"ffmpeg: {p.stderr[-400:]}")
    return p.stdout


def duration(path):
    out = run(["ffprobe", "-v", "error", "-show_entries", "format=duration",
               "-of", "default=nw=1:nk=1", path])
    return float(out.strip())


def download(url, dest):
    with requests.get(url, stream=True, timeout=60) as r:
        r.raise_for_status()
        with open(dest, "wb") as f:
            for chunk in r.iter_content(1 << 20):
                f.write(chunk)


def transcribe(audio_path, source_lang):
    """Retourne ([(debut, fin, texte), ...], langue_detectee)"""
    lang = None if source_lang == "auto" else source_lang
    segments, info = get_whisper().transcribe(audio_path, language=lang, vad_filter=True)
    segs = [(s.start, s.end, s.text.strip()) for s in segments if s.text.strip()]
    return segs, info.language


def translate(texts, src_code, tgt_code, dev="cpu"):
    tok, mdl = get_translator()
    mdl.to(dev)
    tok.src_lang = src_code
    out = []
    for i in range(0, len(texts), 8):
        batch = texts[i:i + 8]
        enc = tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=256).to(dev)
        gen = mdl.generate(**enc, forced_bos_token_id=tok.convert_tokens_to_ids(tgt_code), max_length=256)
        out += tok.batch_decode(gen, skip_special_tokens=True)
    return out


def synthesize(text, model_id, wav_path, dev="cpu"):
    import torch
    import scipy.io.wavfile
    tok, mdl = get_tts(model_id)
    mdl.to(dev)
    inputs = tok(text, return_tensors="pt").to(dev)
    with torch.no_grad():
        wav = mdl(**inputs).waveform[0].cpu().numpy()
    scipy.io.wavfile.write(wav_path, mdl.config.sampling_rate, wav)


def ai_stage(texts, src_code, tgt_code, tts_model, raw_paths, dev="cpu"):
    """Traduction + voix (partie lourde)."""
    translated = translate(texts, src_code, tgt_code, dev)
    for text, path in zip(translated, raw_paths):
        synthesize(text, tts_model, path, dev)
    return translated


# ZeroGPU (Hugging Face) : la partie lourde passe sur la carte graphique gratuite
try:
    import spaces

    @spaces.GPU(duration=90)
    def ai_stage_gpu(texts, src_code, tgt_code, tts_model, raw_paths):
        return ai_stage(texts, src_code, tgt_code, tts_model, raw_paths, "cuda")
except ImportError:  # hors Hugging Face
    def ai_stage_gpu(*args):
        return ai_stage(*args, "cpu")


def preload():
    """Télécharge les modèles au démarrage (hors temps GPU)."""
    get_whisper()
    get_translator()
    for code in ("fon", "yor"):
        get_tts(TARGET_LANGS[code]["tts"])


def fit_to_slot(wav_in, wav_out, slot):
    """Accélère la phrase doublée si elle dépasse le créneau de la phrase d'origine."""
    d = duration(wav_in)
    speed = max(1.0, min(d / max(slot, 0.3), 1.8))  # jamais plus de 1.8x
    run(["ffmpeg", "-y", "-i", wav_in, "-filter:a", f"atempo={speed:.3f}",
         "-ar", "24000", "-ac", "1", wav_out])


def mix(video_in, clips, video_out, total):
    """clips = [(debut_en_s, chemin_wav)]. Garde le son d'origine très bas en fond."""
    cmd = ["ffmpeg", "-y", "-i", video_in]
    for _, p in clips:
        cmd += ["-i", p]
    parts = [f"[0:a]volume={ORIGINAL_VOLUME}[bg]"]
    labels = ["[bg]"]
    for i, (start, _) in enumerate(clips, start=1):
        ms = int(start * 1000)
        parts.append(f"[{i}:a]adelay={ms}|{ms}[c{i}]")
        labels.append(f"[c{i}]")
    parts.append("".join(labels) + f"amix=inputs={len(labels)}:normalize=0:duration=first[aout]")
    cmd += ["-filter_complex", ";".join(parts), "-map", "0:v", "-map", "[aout]",
            "-c:v", "copy", "-c:a", "aac", "-b:a", "128k", "-t", f"{total:.2f}", video_out]
    run(cmd)


# ---------------------------------------------------------------- jobs
jobs = {}
executor = ThreadPoolExecutor(max_workers=1)


def process(job_id, video_url, source_lang, target_lang):
    job = jobs[job_id]
    work = tempfile.mkdtemp(prefix="noukiko_")
    try:
        job.update(status="en_cours", progress=5, step="Téléchargement")
        src = os.path.join(work, "in.mp4")
        download(video_url, src)
        total = duration(src)
        if total > MAX_SECONDS:
            raise RuntimeError(f"Vidéo trop longue ({int(total)} s, max {MAX_SECONDS} s)")

        job.update(progress=15, step="Transcription")
        audio = os.path.join(work, "audio.wav")
        run(["ffmpeg", "-y", "-i", src, "-vn", "-ac", "1", "-ar", "16000", audio])
        segs, detected = transcribe(audio, source_lang)
        if not segs:
            raise RuntimeError("Aucune parole détectée dans la vidéo")
        if detected not in SOURCE_LANGS:
            raise RuntimeError(f"Langue de la vidéo non prise en charge : {detected} (anglais ou français)")
        job["detected_lang"] = detected

        job.update(progress=40, step="Traduction et voix")
        tgt = TARGET_LANGS[target_lang]
        raws = [os.path.join(work, f"raw{i}.wav") for i in range(len(segs))]
        args = ([s[2] for s in segs], SOURCE_LANGS[detected], tgt["nllb"], tgt["tts"], raws)
        try:
            texts = ai_stage_gpu(*args)
        except Exception:  # GPU indisponible -> on continue sur CPU
            job.update(step="Traduction et voix (mode lent)")
            texts = ai_stage(*args, "cpu")
        job["transcript"] = [{"start": round(s[0], 2), "end": round(s[1], 2),
                              "original": s[2], "traduction": t} for s, t in zip(segs, texts)]

        job.update(progress=80, step="Calage des voix")
        clips = []
        for i, (start, end, _) in enumerate(segs):
            fit = os.path.join(work, f"seg{i}.wav")
            fit_to_slot(raws[i], fit, end - start)
            clips.append((start, fit))

        job.update(progress=92, step="Montage")
        out = os.path.join(OUT_DIR, f"{job_id}.mp4")
        mix(src, clips, out, total)
        job.update(status="termine", progress=100,
                   step="Terminé ! Ta vidéo doublée est prête, appuie sur lecture.",
                   result_url=f"{PUBLIC_URL}/files/{job_id}.mp4")
    except Exception as e:  # noqa
        job.update(status="erreur", error=str(e), step=f"Erreur : {e}")
    finally:
        shutil.rmtree(work, ignore_errors=True)


# ---------------------------------------------------------------- API
app = FastAPI(title="Noukiko Dubbing API")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])


class DubRequest(BaseModel):
    video_url: str
    source_lang: str = "auto"
    target_lang: str = "fon"


@app.get("/status")
def root():
    return {"service": "Noukiko", "ok": True}


@app.get("/languages")
def languages():
    return {"sources": list(SOURCE_LANGS),
            "cibles": [{"code": k, "nom": v["name"]} for k, v in TARGET_LANGS.items()]}


@app.post("/dub")
def dub(req: DubRequest):
    if req.source_lang != "auto" and req.source_lang not in SOURCE_LANGS:
        raise HTTPException(400, f"Langue source inconnue : {req.source_lang}")
    if req.target_lang not in TARGET_LANGS:
        raise HTTPException(400, f"Langue cible inconnue : {req.target_lang}")
    if not re.match(r"^https?://", req.video_url):
        raise HTTPException(400, "Le lien doit commencer par http:// ou https://")
    job_id = uuid.uuid4().hex[:12]
    jobs[job_id] = {"status": "en_attente", "progress": 0, "step": "En attente",
                    "result_url": "", "error": ""}
    executor.submit(process, job_id, req.video_url, req.source_lang, req.target_lang)
    return {"job_id": job_id}


@app.get("/jobs/{job_id}")
def job_status(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Tâche introuvable")
    return {"job_id": job_id, **jobs[job_id]}


@app.get("/files/{name}")
def files(name: str):
    path = os.path.join(OUT_DIR, os.path.basename(name))
    if not os.path.exists(path):
        raise HTTPException(404, "Fichier introuvable")
    return FileResponse(path, media_type="video/mp4")


# ---------------------------------------------------------------- page de test (Space Gradio gratuit)
import gradio as gr


def _ui_dub(url, lang):
    r = dub(DubRequest(video_url=url, target_lang=lang))
    return f"Tâche lancée : {r['job_id']} — suivez /jobs/{r['job_id']}"


with gr.Blocks(title="Noukiko") as demo:
    gr.Markdown("# Noukiko – doublage automatique\nAPI : `POST /dub`, `GET /jobs/{id}`")
    u = gr.Textbox(label="Lien direct de la vidéo (.mp4)")
    l = gr.Dropdown(list(TARGET_LANGS), value="fon", label="Langue cible")
    o = gr.Textbox(label="Résultat")
    gr.Button("Doubler").click(_ui_dub, [u, l], o)

# Routes de l'API placées avant celles de Gradio (lancement via demo.launch, requis par ZeroGPU)
_api_routes = [r for r in app.routes if getattr(r, "path", "") in
               ("/status", "/languages", "/dub", "/jobs/{job_id}", "/files/{name}")]

if __name__ == "__main__":
    if os.environ.get("SPACE_ID"):
        preload()
    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", "7860")),
                        app_kwargs={"routes": _api_routes})