from __future__ import annotations

import argparse
import json
import re
import subprocess
from pathlib import Path


def _run(cmd: list[str]) -> None:
    subprocess.run(cmd, check=True)


def _probe_duration_seconds(path: Path) -> float:
    out = subprocess.check_output(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=nw=1:nk=1",
            str(path),
        ],
        text=True,
    ).strip()
    return float(out)


def _fmt_time(seconds: float) -> str:
    ms = int(round(seconds * 1000.0))
    hh = ms // 3_600_000
    ms -= hh * 3_600_000
    mm = ms // 60_000
    ms -= mm * 60_000
    ss = ms // 1_000
    ms -= ss * 1_000
    return f"{hh:02d}:{mm:02d}:{ss:02d},{ms:03d}"


def _silence_path(work_dir: Path, seconds: float) -> Path:
    ms = int(round(seconds * 1000.0))
    return work_dir / f"silence_{ms}ms.wav"


def _ensure_silence(work_dir: Path, seconds: float) -> Path:
    path = _silence_path(work_dir, seconds)
    if path.exists():
        return path
    _run(
        [
            "ffmpeg",
            "-y",
            "-f",
            "lavfi",
            "-i",
            "anullsrc=r=48000:cl=stereo",
            "-t",
            f"{seconds:.3f}",
            "-c:a",
            "pcm_s16le",
            str(path),
        ]
    )
    return path


def _norm_wav(in_wav: Path, out_wav: Path) -> None:
    if out_wav.exists():
        return
    _run(
        [
            "ffmpeg",
            "-y",
            "-i",
            str(in_wav),
            "-ar",
            "48000",
            "-ac",
            "2",
            "-c:a",
            "pcm_s16le",
            str(out_wav),
        ]
    )


def _pause_for(text: str) -> float:
    t = text.strip()
    if t.endswith(("?", "…")):
        return 0.60
    if t.endswith(("!", ".")) and len(t) > 60:
        return 0.45
    if t.endswith((".", "!", ";", ":")):
        return 0.35
    return 0.30


def _safe_stem(s: str) -> str:
    s = s.strip().lower()
    s = re.sub(r"[^a-z0-9]+", "_", s).strip("_")
    return s[:40] or "line"


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Render a Picard/Data conversation using qwen3-tts voice-clone.")
    p.add_argument("--dialogue", required=True, help="Path to dialogue.json")
    p.add_argument("--out-wav", required=True, help="Output wav path")
    p.add_argument("--out-srt", required=True, help="Output srt path")
    p.add_argument("--work-dir", required=True, help="Project work dir (stores clips)")
    p.add_argument("--device-map", default="cpu")
    p.add_argument("--dtype", default="float32")
    return p.parse_args()


def main() -> int:
    args = _parse_args()
    dialogue_path = Path(args.dialogue).expanduser().resolve()
    work_dir = Path(args.work_dir).expanduser().resolve()
    out_wav = Path(args.out_wav).expanduser().resolve()
    out_srt = Path(args.out_srt).expanduser().resolve()

    clips_raw = work_dir / "clips_raw"
    clips_norm = work_dir / "clips_norm"
    clips_raw.mkdir(parents=True, exist_ok=True)
    clips_norm.mkdir(parents=True, exist_ok=True)
    out_wav.parent.mkdir(parents=True, exist_ok=True)
    out_srt.parent.mkdir(parents=True, exist_ok=True)

    root_dir = Path(__file__).resolve().parents[3]
    tts_dir = root_dir / "services" / "TTS" / "implementations" / "qwen3-tts"
    tts_cli = tts_dir / "tts.py"
    tts_venv_py = tts_dir / "runtime" / ".venv" / "bin" / "python"

    items = json.loads(dialogue_path.read_text(encoding="utf-8"))
    if not isinstance(items, list) or not items:
        raise SystemExit("dialogue.json must be a non-empty list")

    # Ensure runtime exists (fast no-op if already installed)
    _run(["python", str(tts_cli), "--install"])

    # Generate all raw clips in one model load (much faster than invoking the CLI per line)
    _run(
        [
            str(tts_venv_py),
            str((work_dir / "render_in_venv.py").resolve()),
            "--dialogue",
            str(dialogue_path),
            "--out-dir",
            str(clips_raw),
            "--device-map",
            args.device_map,
            "--dtype",
            args.dtype,
        ]
    )

    silence_cache: dict[float, Path] = {}

    concat_list = work_dir / "concat.txt"
    srt_lines: list[str] = []
    t_cursor = 0.0
    idx = 1

    concat_entries: list[Path] = []

    for i, item in enumerate(items, start=1):
        speaker = str(item.get("speaker", "")).strip().upper()
        text = str(item.get("text", "")).strip()
        if not speaker or not text:
            raise SystemExit(f"Invalid dialogue entry at index {i}: {item!r}")

        stem = f"{i:03d}_{speaker.lower()}_{_safe_stem(text)}"
        raw_wav = (clips_raw / f"{i:03d}_{speaker.lower()}.wav").resolve()
        norm_wav = (clips_norm / f"{stem}.wav").resolve()

        if not raw_wav.exists():
            raise SystemExit(f"Missing rendered clip for line {i} ({speaker}): {raw_wav}")

        _norm_wav(raw_wav, norm_wav)
        dur = _probe_duration_seconds(norm_wav)

        # Subtitle entry for this line
        start = t_cursor
        end = t_cursor + dur
        srt_lines.append(str(idx))
        srt_lines.append(f"{_fmt_time(start)} --> {_fmt_time(end)}")
        srt_lines.append(f"{speaker}: {text}")
        srt_lines.append("")
        idx += 1

        concat_entries.append(norm_wav)

        pause = _pause_for(text)
        if pause not in silence_cache:
            silence_cache[pause] = _ensure_silence(work_dir, pause)
        concat_entries.append(silence_cache[pause])
        t_cursor = end + pause

    concat_list.write_text("\n".join([f"file '{p.as_posix()}'" for p in concat_entries]) + "\n", encoding="utf-8")

    # Concatenate into final wav
    tmp_out = out_wav.with_suffix(".tmp.wav")
    _run(["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list), "-c", "copy", str(tmp_out)])
    tmp_out.replace(out_wav)

    out_srt.write_text("\n".join(srt_lines), encoding="utf-8")
    return 0


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