#!/usr/bin/env python3
"""Higgs Audio v3 Dubbing — minimal one-click UI.

Upload → pick target language → Generate. Everything else is automatic:
ffmpeg extract → Demucs → Whisper large-v3-turbo → pyannote diarization
→ per-speaker reference extraction → Index-Translate-2B translation
→ per-speaker Higgs Audio v3 TTS + N-GAN 48 kHz → exact-timing mix
→ mux MP4 / export WAV → SRT subtitles.

ZeroGPU Spaces: the single @spaces.GPU marker below is required. On
ZeroGPU torch.cuda.is_available() is False inside non-@spaces.GPU code,
so Demucs/translation/TTS automatically stay on CPU. On a normal CUDA
host they use GPU.
"""
from __future__ import annotations

import os


# ---------------------------------------------------------------------------
# 0. CPU detection before torch/numpy.
# ---------------------------------------------------------------------------
def _cgroup_cpu_quota() -> float | None:
    try:
        with open("/sys/fs/cgroup/cpu.max") as f:
            parts = f.read().split()
        if len(parts) == 2 and parts[0] != "max":
            q, p = int(parts[0]), int(parts[1])
            if p > 0:
                return q / p
    except (FileNotFoundError, ValueError, OSError):
        pass
    try:
        with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as f:
            q = int(f.read().strip())
        with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as f:
            p = int(f.read().strip())
        if q > 0 and p > 0:
            return q / p
    except (FileNotFoundError, ValueError, OSError):
        pass
    return None


def _available_cpus() -> int:
    try:
        affinity = len(os.sched_getaffinity(0))
    except AttributeError:
        affinity = os.cpu_count() or 1
    quota = _cgroup_cpu_quota()
    if quota is not None:
        return min(affinity, max(1, int(round(quota))))
    return max(1, affinity)


_CPU_COUNT = _available_cpus()
os.environ.setdefault("OMP_NUM_THREADS", str(_CPU_COUNT))
os.environ.setdefault("MKL_NUM_THREADS", str(_CPU_COUNT))
os.environ.setdefault("OPENBLAS_NUM_THREADS", str(_CPU_COUNT))
os.environ.setdefault("NUMEXPR_NUM_THREADS", str(_CPU_COUNT))
os.environ.setdefault("VECLIB_MAXIMUM_THREADS", str(_CPU_COUNT))
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")

try:
    _aff = len(os.sched_getaffinity(0))
except AttributeError:
    _aff = os.cpu_count() or 1
print(
    f"[cpu] cpu_count={os.cpu_count()}  sched_getaffinity={_aff}  "
    f"cgroup_quota={_cgroup_cpu_quota()}  -> using {_CPU_COUNT} threads  "
    f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')!r})",
    flush=True,
)


# ---------------------------------------------------------------------------
# 1. spaces shim + ZeroGPU startup stub.
# ---------------------------------------------------------------------------
try:
    import spaces  # noqa: F401
except ImportError:
    class _SpacesShim:
        def GPU(self, *args, **kwargs):
            if len(args) == 1 and callable(args[0]) and not kwargs:
                return args[0]
            def deco(fn):
                return fn
            return deco
        def __getattr__(self, name):
            def _noop(*a, **kw):
                if len(a) == 1 and callable(a[0]) and not kw:
                    return a[0]
                return lambda fn: fn
            return _noop
    spaces = _SpacesShim()


@spaces.GPU(duration=10)
def _zerogpu_startup_probe():
    """Marker for ZeroGPU; never invoked."""
    return True


# ---------------------------------------------------------------------------
# 2. Regular imports.
# ---------------------------------------------------------------------------
import argparse
import logging
import shutil
import time
from pathlib import Path
from typing import Optional, Union

import gradio as gr
import numpy as np
import pandas as pd
import soundfile as sf

try:
    import torch
    torch.set_num_threads(_CPU_COUNT)
    try:
        torch.set_num_interop_threads(max(1, _CPU_COUNT // 2))
    except RuntimeError:
        pass
except ImportError:
    pass

from backend.core.state import STORE, Job, Segment
from backend.services import (
    asr, diarization, media, persistence, separation, subtitles, translation,
)
from backend.services.higgs_tts import HiggsTTSService


# ---------------------------------------------------------------------------
# 3. Paths + logging.
# ---------------------------------------------------------------------------
DATA_DIR = Path(os.environ.get("DUB_DATA", "data/jobs"))
PROFILES_DIR = Path(os.environ.get("DUB_PROFILES", "profiles"))
LOG_DIR = Path(os.environ.get("DUB_LOGS", "data/logs"))

DATA_DIR.mkdir(parents=True, exist_ok=True)
PROFILES_DIR.mkdir(parents=True, exist_ok=True)
LOG_DIR.mkdir(parents=True, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(LOG_DIR / "app.log", encoding="utf-8"),
    ],
)
logger = logging.getLogger("app")

_SESSION_PREFIX = "_session_"

CSS = """
.gradio-container {max-width: 100% !important; font-size: 15px !important;}
.log-box textarea {font-family: ui-monospace, Menlo, Consolas, monospace;
                   font-size: 12px !important;}
"""

LANG_CODES = [
    "en","de","fr","es","it","pt","ru","ja","ko","zh","ar","hi",
    "tr","pl","nl","uk","yue","vi","th","id","cs","sv","da","fi",
    "no","el","he","fa",
]
LANG_LABELS = {
    "en":"English","de":"German","fr":"French","es":"Spanish","it":"Italian",
    "pt":"Portuguese","ru":"Russian","ja":"Japanese","ko":"Korean","zh":"Chinese",
    "ar":"Arabic","hi":"Hindi","tr":"Turkish","pl":"Polish","nl":"Dutch",
    "uk":"Ukrainian","yue":"Cantonese","vi":"Vietnamese","th":"Thai",
    "id":"Indonesian","cs":"Czech","sv":"Swedish","da":"Danish","fi":"Finnish",
    "no":"Norwegian","el":"Greek","he":"Hebrew","fa":"Persian",
}


# ---------------------------------------------------------------------------
# 4. Helpers.
# ---------------------------------------------------------------------------
def _resolve_path(x: Union[str, Path, object, None]) -> Optional[str]:
    if x is None:
        return None
    if isinstance(x, (str, Path)):
        return str(x)
    for a in ("path", "name"):
        v = getattr(x, a, None)
        if isinstance(v, (str, Path)):
            return str(v)
    return str(x)


def _best_segment(segs: list[dict]) -> dict:
    def score(s):
        d = s["end"] - s["start"]
        return (2, d) if 3.0 <= d <= 10.0 else (1, d)
    return max(segs, key=score)


def _extract_speaker_refs(
    audio_path: Path, segments: list[dict], work: Path,
    profiles: Path, job_id: str, log,
) -> dict[str, dict]:
    by_spk: dict[str, list[dict]] = {}
    for s in segments:
        spk = s.get("speaker") or ""
        if spk:
            by_spk.setdefault(spk, []).append(s)
    if not by_spk:
        return {}
    full, sr = sf.read(str(audio_path), dtype="float32", always_2d=False)
    if full.ndim > 1:
        full = full.mean(axis=1)
    out_dir = work / "speakers"
    out_dir.mkdir(parents=True, exist_ok=True)
    job_short = job_id[-8:] if len(job_id) > 8 else job_id
    refs: dict[str, dict] = {}
    for spk, segs in sorted(by_spk.items()):
        best = _best_segment(segs)
        start = max(0, int(best["start"] * sr))
        end = min(len(full), int(best["end"] * sr))
        chunk = full[start:end][: int(10.0 * sr)]
        if len(chunk) < int(1.0 * sr):
            log(f"  ⚠️  {spk}: reference too short — skipped")
            continue
        ref_path = out_dir / f"{spk}.wav"
        sf.write(str(ref_path), chunk, sr)
        session_name = f"{_SESSION_PREFIX}{job_short}_{spk}"
        sf.write(str(profiles / f"{session_name}.wav"), chunk, sr)
        refs[spk] = {
            "path": str(ref_path),
            "ref_text": (best.get("text") or "").strip() or None,
            "session_name": session_name,
            "duration": len(chunk) / sr,
        }
        log(f"  ✓ {spk}: {len(chunk) / sr:.1f}s → {session_name}")
    return refs


def _fit_exact(audio: np.ndarray, sr: int, target_sec: float) -> np.ndarray:
    if target_sec <= 0:
        return audio
    target = int(target_sec * sr)
    if len(audio) > target:
        return audio[:target]
    if len(audio) < target:
        return np.pad(audio, (0, target - len(audio)))
    return audio


def _resolve_reference(seg: Segment, job: Job) -> tuple[Optional[str], Optional[str]]:
    if seg.speaker and seg.speaker in job.speaker_refs:
        info = job.speaker_refs[seg.speaker]
        return info["path"], info["ref_text"]
    if job.global_ref_audio:
        return job.global_ref_audio, None
    return None, None


def _scan_first_profile() -> Optional[str]:
    """Pick the first permanent profile as global fallback, if any."""
    for p in sorted(PROFILES_DIR.glob("*.wav")):
        if not p.stem.startswith(_SESSION_PREFIX):
            return p.stem
    return None


# ---------------------------------------------------------------------------
# 5. One-click pipeline.
# ---------------------------------------------------------------------------
def dub(
    media_path,
    target_lang: str,
    progress: gr.Progress = gr.Progress(),
):
    log_lines: list[str] = []

    def log(msg: str) -> str:
        log_lines.append(msg)
        logger.info(msg)
        return "\n".join(log_lines)

    empty = (None, None, None, None, "")
    path = _resolve_path(media_path)
    if not path or not Path(path).exists():
        return (*empty, "❌ Upload a video or audio file first.")

    # ---- 1. Prep ----------------------------------------------------
    src = Path(path)
    media_type = media.classify_media(src)
    job_id = f"dub_{time.strftime('%Y%m%d_%H%M%S')}"
    work = DATA_DIR / job_id
    work.mkdir(parents=True, exist_ok=True)
    src_copy = work / src.name
    shutil.copyfile(src, src_copy)

    job = STORE.create(work, src_copy, media_type)
    job.target_lang = target_lang
    log(f"→ Job {job.job_id}  ({media_type})")

    progress(0.05, desc="Extracting audio")
    log("→ Extracting audio (48 kHz mono) ...")
    audio_wav = work / "original.wav"
    media.extract_audio(src_copy, audio_wav, sr=48000)
    job.audio_path = audio_wav
    job.duration = media.get_duration(src_copy)
    log(f"  ✓ {job.duration:.1f}s")

    # ---- 2. Demucs --------------------------------------------------
    progress(0.12, desc="Separating vocals")
    log("→ Demucs separation ...")
    if separation.available():
        try:
            _v, bg = separation.separate(audio_wav, work / "demucs")
            job.vocals_path, job.background_path = _v, bg
            log("  ✓ vocals + background")
        except Exception as e:
            log(f"  ⚠️  Demucs failed: {e} — silence background")
            silent = work / "silence.wav"
            sf.write(str(silent), np.zeros(int(job.duration * 48000),
                                            dtype=np.float32), 48000)
            job.background_path = silent
    else:
        log("  ⚠️  Demucs not installed — original audio as background")
        job.background_path = audio_wav

    # ---- 3. ASR -----------------------------------------------------
    progress(0.22, desc="Transcribing")
    log("→ Whisper large-v3-turbo ...")
    segs, src_lang = asr.transcribe(audio_wav)
    job.source_lang = src_lang
    log(f"  ✓ {len(segs)} segments, source={src_lang}")

    # ---- 4. Diarization --------------------------------------------
    job.speaker_refs = {}
    if diarization.available():
        progress(0.32, desc="Diarizing")
        log("→ Diarization (auto-detect speakers) ...")
        try:
            turns = diarization.diarize(audio_wav, num_speakers=None)
            segs = diarization.assign_speakers(segs, turns)
            unique = sorted({s.get("speaker") for s in segs if s.get("speaker")})
            log(f"  ✓ {len(unique)} speaker(s): {unique}")
        except Exception as e:
            log(f"  ⚠️  Diarization failed: {e}")

    if any(s.get("speaker") for s in segs):
        progress(0.40, desc="Cloning speakers")
        log("→ Extracting per-speaker clones ...")
        job.speaker_refs = _extract_speaker_refs(
            audio_wav,
            [{"start": s["start"], "end": s["end"],
              "speaker": s.get("speaker", ""), "text": s["text"]} for s in segs],
            work, PROFILES_DIR, job_id, log,
        )

    # ---- 5. Translate ----------------------------------------------
    progress(0.50, desc="Translating")
    log(f"→ Translating → {target_lang} ...")
    translated = translation.translate_segments(
        [{"id": i, "text": s["text"]} for i, s in enumerate(segs)],
        target_lang, None,
    )
    by_id = {t["id"]: t["text"] for t in translated}
    log(f"  ✓ {len(translated)} translated")

    job.segments = [
        Segment(
            idx=i,
            start=float(s["start"]),
            end=float(s["end"]),
            text_original=s["text"],
            text_translated=by_id.get(i, s["text"]),
            speaker=s.get("speaker", ""),
        )
        for i, s in enumerate(segs)
    ]

    # ---- 6. Global fallback (first permanent profile, if any) ------
    job.global_ref_audio = ""
    fallback = _scan_first_profile()
    if fallback:
        job.global_ref_audio = str(PROFILES_DIR / f"{fallback}.wav")
        log(f"  Fallback voice: {fallback}")

    # ---- 7. TTS per segment ----------------------------------------
    progress(0.55, desc="Loading TTS")
    log("→ Loading Higgs Audio v3 + N-GAN ...")
    tts = HiggsTTSService.get()
    tts.load()
    sr = tts.sampling_rate
    log(f"  ✓ TTS ready (sr={sr})")

    seg_dir = work / "segments"
    seg_dir.mkdir(exist_ok=True)
    n_ok = n_fail = n_skip = 0

    log(f"→ Synthesizing {len(job.segments)} segment(s) ...")
    for i, seg in enumerate(job.segments):
        text = (seg.text_translated or seg.text_original or "").strip()
        if not text:
            n_skip += 1
            continue
        ref_audio, ref_text = _resolve_reference(seg, job)
        if ref_audio is None:
            n_skip += 1
            log(f"  ⚠️  seg {seg.idx}: no reference voice — skipped")
            continue
        try:
            _, arr = tts.generate(
                text,
                reference_audio=ref_audio,
                reference_text=ref_text,
                temperature=0.8,
                top_p=0.95,
                max_new_tokens=2048,
                seed=-1,
            )
            audio_f = arr.astype(np.float32) / 32768.0
            audio_f = media.fit_to_duration(audio_f, sr, seg.duration)
            out_path = seg_dir / f"seg_{seg.idx:05d}.wav"
            sf.write(str(out_path), audio_f, sr)
            seg.audio_path = str(out_path)
            n_ok += 1
        except Exception as e:
            log(f"  ❌ seg {seg.idx} ({seg.speaker or '-'}): "
                f"{type(e).__name__}: {e}")
            n_fail += 1

        progress(0.55 + 0.30 * (i + 1) / len(job.segments),
                 desc=f"TTS {i + 1}/{len(job.segments)}")

    log(f"  ✓ {n_ok} ok, {n_fail} failed, {n_skip} skipped")
    if n_ok == 0:
        return (*empty, "❌ Nothing synthesized — no reference voices available.")

    # ---- 8. Assemble + mix -----------------------------------------
    progress(0.88, desc="Mixing")
    log("→ Assembling + mixing ...")
    voice_track = work / "voice.wav"
    segs_with_audio = [s for s in job.segments if s.audio_path]
    max_end = max((s.start + s.duration for s in segs_with_audio),
                  default=job.duration)
    media.assemble_voice_track(
        [{"start": s.start, "audio_path": s.audio_path} for s in segs_with_audio],
        max(job.duration, max_end) + 1.0, sr, voice_track,
    )
    job.voice_track_path = voice_track

    mixed = work / "mixed.wav"
    media.mix_tracks(
        voice=voice_track, background=job.background_path,
        sr=sr, out_path=mixed, voice_gain=1.0, bg_gain=0.35,
    )
    job.mixed_path = mixed
    log(f"  ✓ {mixed.name}")

    # ---- 9. Export -------------------------------------------------
    progress(0.95, desc="Exporting")
    if job.media_type == "video":
        final = work / "output.mp4"
        log("→ Muxing MP4 ...")
        media.mux_video(job.media_path, mixed, final)
        job.output_path = final
        video_out, audio_out = str(final), None
    else:
        final = work / "output.wav"
        log("→ Exporting WAV ...")
        media.export_audio(mixed, final, codec="wav")
        job.output_path = final
        video_out, audio_out = None, str(final)

    sub_path = work / "subtitles.srt"
    sub_data = [
        {"start": s.start, "end": s.end,
         "text_original": s.text_original,
         "text_translated": s.text_translated}
        for s in job.segments
    ]
    subtitles.export_srt(sub_data, sub_path, use_translation=True, offset_s=0.0)
    log(f"  ✓ subtitles.srt")

    persistence.save_job(job)
    progress(1.0, desc="Done")
    log(f"✓ Done: {job.output_path}")
    return video_out, audio_out, str(final), str(sub_path), "\n".join(log_lines)


# ---------------------------------------------------------------------------
# 6. UI.
# ---------------------------------------------------------------------------
def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Dubbing") as demo:
        gr.Markdown(
            "# 🎬 Video & Audio Dubbing\n"
            "Upload → pick target language → **Generate**. "
            "Everything else is automatic."
        )

        with gr.Row():
            with gr.Column(scale=1):
                media_in = gr.File(
                    label="Input (video or audio)",
                    file_types=["video", "audio"],
                )
                target_lang = gr.Dropdown(
                    choices=[(LANG_LABELS[c], c) for c in LANG_CODES],
                    value="ru",
                    label="Target language",
                )
                run_btn = gr.Button(
                    "🎙️ Generate", variant="primary", size="lg",
                )
                log_box = gr.Textbox(
                    label="Log", lines=22, max_lines=60,
                    interactive=False, elem_classes="log-box",
                )

            with gr.Column(scale=1):
                video_out = gr.Video(label="Dubbed video")
                audio_out = gr.Audio(label="Dubbed audio")
                download = gr.File(label="Download")
                subtitle_out = gr.File(label="Subtitles")

        run_btn.click(
            dub,
            inputs=[media_in, target_lang],
            outputs=[video_out, audio_out, download, subtitle_out, log_box],
        )

        gr.Markdown(
            "<sub>All processing runs automatically: Demucs, Whisper, "
            "diarization, translation, per-speaker cloning, TTS, mix, "
            "mux, subtitles. Voices auto-cloned from each detected speaker.</sub>"
        )

    return demo


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=7860)
    ap.add_argument("--ip", default="0.0.0.0")
    ap.add_argument("--share", action="store_true")
    args = ap.parse_args()

    demo = build_ui()
    demo.queue(default_concurrency_limit=1).launch(
        server_name=args.ip,
        server_port=args.port,
        share=args.share,
        theme=gr.themes.Soft(),
        css=CSS,
        ssr_mode=False,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())