#!/usr/bin/env python3
"""Ad localizer helper. Two sub-commands:

  plan     print the Kenerate MCP calls + output names for each language (runs nothing, costs nothing)
  collect  download every locale's dubbed / captioned result, name it <name>_<lang>.mp4, check duration and audio
           against the source, optionally burn a translated SRT locally (free alternative to add_subtitles),
           and write locales.json / locales.csv + a contact sheet.

Examples
  python localize.py plan --source https://.../ad.mp4 --langs es,fr,de,hi,ja --captions
  python localize.py collect --name ad --source https://.../ad.mp4 --item es=https://.../dub_es.mp4 \
      --item fr=https://.../dub_fr.mp4 --srt fr=ad_fr.srt
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import kmedia as K  # noqa: E402

# dub_video target_lang -> add_subtitles language (only the locales add_subtitles lists; others auto-detect)
SUB_LANG = {"en": "en-US", "es": "es-ES", "fr": "fr-FR", "de": "de-DE", "ja": "ja-JP", "zh": "zh-CN", "it": "it-IT",
            "pt": "pt-BR", "ru": "ru-RU", "hi": "", "ar": "", "ko": "", "id": "", "tr": "", "nl": "", "pl": ""}


def plan(a) -> None:
    langs = [l.strip() for l in a.langs.split(",") if l.strip()]
    name = K.safe_name(a.name)
    steps = []
    for l in langs:
        s = {"lang": l, "output": f"{name}_{l}.mp4",
             "calls": [{"tool": "estimate_cost", "arguments": {"tool": "dub_video", "arguments": {"media_url": a.source, "target_lang": l}}},
                       {"tool": "dub_video", "arguments": {"media_url": a.source, "target_lang": l,
                                                           **({"source_lang": a.source_lang} if a.source_lang else {})}},
                       {"tool": "wait_for_generation", "arguments": {"kind": "dubbing", "id": "<job id>"}}]}
        if a.captions:
            sub = {"video_url": "<dubbed url>", "style": a.style}
            if SUB_LANG.get(l):
                sub["language"] = SUB_LANG[l]
            s["calls"] += [{"tool": "add_subtitles", "arguments": sub},
                           {"tool": "wait_for_generation", "arguments": {"kind": "subtitles", "id": "<job id>"}}]
        steps.append(s)
    K.finish({"source": a.source, "languages": langs, "plan": steps,
              "then": f"python localize.py collect --name {name} --source <source> " + " ".join(f"--item {l}=<final url>" for l in langs),
              "notes": ["dub_video is priced per second: run estimate_cost for one language and multiply; confirm with the user.",
                        "add_subtitles captions only speech, up to 4 min; 1080p doubles the price.",
                        "Languages add_subtitles does not list (e.g. hi, ar, ko) are auto-detected."]})


def collect(a) -> None:
    items = []
    for it in a.item:
        l, _, src = it.partition("=")
        if not src:
            K.fail(f"Bad --item '{it}'.", "Use LANG=URL_OR_FILE, e.g. --item es=https://...")
        items.append((l.strip(), src.strip()))
    srts = dict((x.partition("=")[0].strip(), x.partition("=")[2].strip()) for x in a.srt)
    st = K.Status(4, "ad-localizer")
    K.require(status=st)
    tmp = K.workdir()
    out = Path(a.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    name = K.safe_name(a.name)

    st.step(f"Downloading {len(items)} locale(s)" + (" + source" if a.source else ""))
    ref = K.probe(K.fetch(a.source, tmp, "source", st)) if a.source else None
    got = []
    for l, src in items:
        p = K.fetch(src, tmp, f"{name}_{l}_raw", st)
        got.append((l, p))

    st.step("Checking duration and audio")
    rows = []
    for l, p in got:
        info = K.probe(p)
        row = {"lang": l, "duration": round(info["duration"], 2), "size": f"{info['width']}x{info['height']}",
               "has_audio": info["has_audio"], "issues": []}
        if not info["has_audio"]:
            row["issues"].append("no audio track")
        if ref and ref["duration"] and abs(info["duration"] - ref["duration"]) / ref["duration"] > a.tolerance:
            row["issues"].append(f"length {info['duration']:.1f}s vs source {ref['duration']:.1f}s")
        rows.append((row, p))
        if row["issues"]:
            st.warn(f"{l}: " + "; ".join(row["issues"]))

    st.step(f"Naming files{' + burning SRTs' if srts else ''}")
    final = []
    for i, (row, p) in enumerate(rows, 1):
        l = row["lang"]
        dst = out / f"{name}_{l}{p.suffix or '.mp4'}"
        if l in srts:
            cues = K.parse_srt(Path(srts[l]).read_text(encoding="utf-8", errors="replace"))
            info = K.probe(p)
            lst = K.caption_track(cues, info["width"], info["height"], info["duration"], tmp / f"caps_{l}", a.caption_style,
                                  a.caption_pos, a.font)
            K.overlay_track(p, lst, dst, st)
            row["captions"] = f"burned locally from {Path(srts[l]).name} ({len(cues)} cues)"
        else:
            import shutil
            shutil.copyfile(p, dst)
        row["file"] = str(dst.resolve())
        final.append(dst)
        st.progress(i * 100 / len(rows), dst.name)

    st.step("Report + contact sheet")
    frames = []
    for row, _ in rows:
        try:
            frames.append(K.extract_frame(Path(row["file"]), row["duration"] / 2, tmp / f"f_{row['lang']}.jpg"))
        except K.FFmpegError:
            pass
    sheet = K.contact_sheet(frames, out / f"{name}_locales.jpg", cols=min(5, len(frames)), cell=300,
                            labels=[r["lang"] for r, _ in rows]) if frames else None
    data = [r for r, _ in rows]
    (out / f"{name}_locales.json").write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
    with open(out / f"{name}_locales.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["lang", "file", "duration", "has_audio", "issues"])
        for r in data:
            w.writerow([r["lang"], r["file"], r["duration"], r["has_audio"], "; ".join(r["issues"])])
    K.finish({"locales": data, "sheet": str(sheet.resolve()) if sheet else None,
              "source_duration": round(ref["duration"], 2) if ref else None, "elapsed_s": st.elapsed()})


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("plan")
    p.add_argument("--source", required=True)
    p.add_argument("--langs", required=True, help="dub_video codes, e.g. es,fr,de,hi,ja")
    p.add_argument("--source-lang", default="")
    p.add_argument("--captions", action="store_true", help="add add_subtitles after each dub")
    p.add_argument("--style", default="simple")
    p.add_argument("--name", default="ad")
    c = sub.add_parser("collect")
    c.add_argument("--item", action="append", default=[], required=True, help="LANG=URL_OR_FILE (repeat)")
    c.add_argument("--source", default="", help="original ad, to compare lengths")
    c.add_argument("--srt", action="append", default=[], help="LANG=file.srt to burn locally")
    c.add_argument("--caption-style", choices=K.CAPTION_STYLES, default="box")
    c.add_argument("--caption-pos", choices=["bottom", "center", "top"], default="bottom")
    K.add_font_args(c)
    c.add_argument("--tolerance", type=float, default=0.08, help="allowed length difference (default 8%%)")
    c.add_argument("--name", default="ad")
    c.add_argument("--out-dir", default="kenerate-out/ad-localizer")
    a = K.parse_args(ap)
    plan(a) if a.cmd == "plan" else collect(a)


if __name__ == "__main__":
    K.main_guard(main)