#!/usr/bin/env python3
"""Word-by-word animated captions (TikTok / Reels style) burned into a video, plus an .ass karaoke file and a
word-timing .json you can reuse in an editor.

Timing sources, best first:
  --words-json  [{"word","start","end"}, ...]  or ElevenLabs alignment {characters, character_start_times_seconds, ...}
  --ass         ASS / SSA with {\\k} karaoke tags (exact word timing) or plain dialogue lines
  --srt         SRT / VTT cues: word times are estimated inside each cue (by word length)
  --text        a script: spread over the whole video (rough)

Styles: highlight (active word on an accent box), karaoke (spoken words fill with the accent), pop (words pop in
as they are said), bounce (one big word at a time), typewriter (letters type on), minimal (clean, active word bright).

Examples
  python captions.py --video talk.mp4 --srt talk.srt --style highlight --accent "#FFD60A" --font-family Montserrat --font-weight 800
  python captions.py --video reel.mp4 --words-json words.json --style pop --max-words 2 --position center
  python captions.py --video reel.mp4 --srt a.srt --style karaoke --preview-only     # contact sheet, no export
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

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

STYLES = ("highlight", "karaoke", "pop", "bounce", "typewriter", "minimal")


# ── timing ───────────────────────────────────────────────────────────────────
def words_from_cues(cues):
    words = []
    for a, b, txt in cues:
        ws = txt.split()
        if not ws:
            continue
        weights = [len(re.sub(r"\W", "", w)) + 2 for w in ws]
        tot, t = sum(weights), a
        for w, k in zip(ws, weights):
            d = (b - a) * k / tot
            words.append({"word": w, "start": round(t, 3), "end": round(t + d, 3), "cue_end": b})
            t += d
    return words


def words_from_json(path: Path):
    data = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(data, dict) and "characters" in data:  # ElevenLabs alignment
        chars, st_, en = data["characters"], data["character_start_times_seconds"], data["character_end_times_seconds"]
        words, cur, s0, e0 = [], "", None, None
        for c, s, e in zip(chars, st_, en):
            if c.isspace():
                if cur:
                    words.append({"word": cur, "start": s0, "end": e0})
                cur, s0 = "", None
                continue
            if s0 is None:
                s0 = s
            cur += c
            e0 = e
        if cur:
            words.append({"word": cur, "start": s0, "end": e0})
        return words
    if isinstance(data, dict):
        data = data.get("words") or data.get("segments") or []
    out = []
    for w in data:
        if "words" in w:  # whisper-style segments
            out += [{"word": x.get("word", x.get("text", "")).strip(), "start": float(x["start"]), "end": float(x["end"])} for x in w["words"]]
        else:
            out.append({"word": str(w.get("word", w.get("text", ""))).strip(), "start": float(w["start"]), "end": float(w["end"])})
    return [w for w in out if w["word"]]


def words_from_ass(path: Path):
    words, cues = [], []
    for ln in path.read_text(encoding="utf-8", errors="replace").splitlines():
        if not ln.startswith("Dialogue:"):
            continue
        parts = ln.split(",", 9)
        if len(parts) < 10:
            continue
        a, b, text = K._ts(parts[1]), K._ts(parts[2]), parts[9]
        ks = re.findall(r"\{[^}]*\\[kK][fo]?(\d+)[^}]*\}([^{]*)", text)
        if ks:
            t = a
            for cs, w in ks:
                d = int(cs) / 100.0
                for piece in w.split():
                    words.append({"word": piece, "start": round(t, 3), "end": round(t + d, 3), "cue_end": b})
                t += d
        else:
            cues.append((a, b, re.sub(r"\{[^}]*\}", "", text).replace("\\N", " ").strip()))
    return words + words_from_cues(cues)


def pages_from_words(words, max_words: int, max_chars: int, gap: float):
    pages, cur = [], []
    for i, w in enumerate(words):
        if cur:
            prev = cur[-1]
            chars = sum(len(x["word"]) + 1 for x in cur) + len(w["word"])
            if (len(cur) >= max_words or chars > max_chars or w["start"] - prev["end"] > gap
                    or re.search(r"[.!?;:]$", prev["word"]) or prev.get("cue_end") != w.get("cue_end")):
                pages.append(cur)
                cur = []
        cur.append(w)
    if cur:
        pages.append(cur)
    out = []
    for i, p in enumerate(pages):
        start = p[0]["start"]
        nxt = pages[i + 1][0]["start"] if i + 1 < len(pages) else None
        end = p[-1]["end"] + 0.25
        if nxt is not None:
            end = min(end, nxt)
        out.append({"words": p, "start": start, "end": max(end, p[-1]["end"])})
    return out


# ── rendering ────────────────────────────────────────────────────────────────
class Renderer:
    def __init__(self, a, W: int, H: int):
        self.a, self.W, self.H = a, W, H
        self.short = min(W, H)
        self.accent = K.hex_rgb(a.accent)
        self.color = K.hex_rgb(a.color)
        self.cache = {}
        vertical = H > W
        self.cy = {"top": 0.2, "center": 0.5, "lower": 0.68 if vertical else 0.8, "bottom": 0.78 if vertical else 0.86}[a.position]
        self.max_w = int(W * (0.84 if vertical else 0.7))
        self.size = int(self.short * 0.1 * a.size) if a.style != "minimal" else int(self.short * 0.062 * a.size)
        if a.style == "bounce":
            self.size = int(self.size * 1.5)
        self.upper = a.style != "minimal" and not a.keep_case
        self.on_accent = K.best_text_color(self.accent)[0]

    def txt(self, w: str):
        w = w.upper() if self.upper else w
        return re.sub(r"\s+", " ", w)

    def sprite(self, word: str, color, size: int, stroke=True):
        key = (word, tuple(color), size, stroke)
        if key not in self.cache:
            st = max(2, int(size * 0.09)) if stroke else 0
            kind = "regular" if self.a.style == "minimal" and not self.a.font else "bold"
            self.cache[key] = K.text_img(word, size, tuple(color) + (255,), kind, self.a.font, stroke=st, stroke_fill=(0, 0, 0, 255))
        return self.cache[key]

    def layout(self, page):
        """Word positions (cx, cy) for a page, shrinking the font until it fits in 2 lines."""
        key = ("layout", id(page))
        if key in self.cache:
            return self.cache[key]
        size = self.size
        while True:
            sprites = [self.sprite(self.txt(w["word"]), self.color, size, self.a.style != "minimal") for w in page["words"]]
            space = int(size * 0.28)
            lines, cur, cw = [], [], 0
            for i, s in enumerate(sprites):
                need = s.width + (space if cur else 0)
                if cur and cw + need > self.max_w:
                    lines.append(cur)
                    cur, cw = [], 0
                    need = s.width
                cur.append(i)
                cw += need
            lines.append(cur)
            if len(lines) <= 2 or size < self.size * 0.55:
                break
            size = int(size * 0.9)
        lh = int(size * 1.18)
        top = self.H * self.cy - lh * len(lines) / 2
        pos = {}
        for li, ln in enumerate(lines):
            lw = sum(sprites[i].width for i in ln) + space * (len(ln) - 1)
            x = self.W / 2 - lw / 2
            for i in ln:
                pos[i] = (x + sprites[i].width / 2, top + li * lh + lh / 2)
                x += sprites[i].width + space
        self.cache[key] = (size, pos)
        return size, pos

    def draw(self, page, k: int, dt: float):
        """Frame for page with word k active, dt seconds after that word started."""
        from PIL import Image, ImageDraw
        img = Image.new("RGBA", (self.W, self.H), (0, 0, 0, 0))
        st = self.a.style
        words = page["words"]
        if st == "bounce":
            w = words[k]
            s = self.sprite(self.txt(w["word"]), self.accent if k % 2 else self.color, self.size)
            p = K.ease_out_back(dt / 0.22, 2.2)
            K.place(img, s, self.W / 2, self.H * self.cy + (1 - K.clamp(dt / 0.22)) * self.size * 0.35, 0.6 + 0.4 * p, K.clamp(dt / 0.08))
            return img
        size, pos = self.layout(page)
        if st == "typewriter":
            text = " ".join(self.txt(w["word"]) for w in words)
            t_abs, parts = page["start"] + dt, []
            for w in words:  # each word types on over its first ~0.25 s
                frac = K.clamp((t_abs - w["start"]) / max(min(0.25, w["end"] - w["start"]), 0.05))
                ww = self.txt(w["word"])
                if frac <= 0:
                    break
                parts.append(ww[:max(1, int(round(len(ww) * frac)))])
                if frac < 1:
                    break
            shown = " ".join(parts) or text[:1]
            spr = self.sprite(shown, self.color, size)
            full = self.sprite(text, self.color, size)
            cx, cy = self.W / 2 - full.width / 2 + spr.width / 2, self.H * self.cy
            if full.width > self.max_w:  # long page: fall back to centred wrap via fit
                spr = K.text_img(shown, int(size * self.max_w / full.width), tuple(self.color) + (255,), "bold", self.a.font,
                                 stroke=max(2, int(size * 0.07)))
                cx = self.W / 2 - full.width * (self.max_w / full.width) / 2 + spr.width / 2
            K.place(img, spr, cx, cy)
            return img
        for i, w in enumerate(words):
            cx, cy = pos[i]
            word = self.txt(w["word"])
            if st == "highlight":
                if i == k:
                    s = self.sprite(word, self.on_accent, size, stroke=False)
                    pad = int(size * 0.16)
                    box = Image.new("RGBA", (s.width + pad * 2, s.height + pad), (0, 0, 0, 0))
                    ImageDraw.Draw(box).rounded_rectangle((0, 0, box.width - 1, box.height - 1), radius=int(size * 0.2),
                                                          fill=self.accent + (255,))
                    sc = 0.85 + 0.15 * K.ease_out_back(dt / 0.14, 2.0)
                    K.place(img, box, cx, cy, sc, 1.0, rot=-2 if (k % 2) else 2)
                    K.place(img, s, cx, cy, sc)
                else:
                    K.place(img, self.sprite(word, self.color, size), cx, cy)
            elif st == "karaoke":
                col = self.accent if i <= k else self.color
                sc = 1.0 + 0.08 * (1 - K.ease_out_cubic(dt / 0.2)) if i == k else 1.0
                K.place(img, self.sprite(word, col, size), cx, cy, sc)
            elif st == "pop":
                if i > k:
                    continue
                s = self.sprite(word, self.accent if i == k else self.color, size)
                sc = 0.55 + 0.45 * K.ease_out_back(dt / 0.18, 2.4) if i == k else 1.0
                K.place(img, s, cx, cy, sc, K.clamp(dt / 0.06) if i == k else 1.0)
            else:  # minimal
                s = self.sprite(word, self.color, size, stroke=False)
                if i != k:
                    s = K.with_opacity(s, 0.55)
                K.place(img, s, cx, cy)
        if st == "minimal":  # soft shadow for legibility
            from PIL import ImageFilter
            sh = Image.new("RGBA", img.size, (0, 0, 0, 0))
            sh.putalpha(img.getchannel("A").point(lambda v: int(v * 0.6)).filter(ImageFilter.GaussianBlur(max(2, size // 12))))
            sh.alpha_composite(img)
            img = sh
        return img


def write_ass(pages, dst: Path, W: int, H: int, a, font_name: str):
    def ts(t):
        cs = int(round(t * 100))
        return f"{cs // 360000}:{cs // 6000 % 60:02d}:{cs // 100 % 60:02d}.{cs % 100:02d}"

    def col(rgb):
        return "&H00{:02X}{:02X}{:02X}".format(rgb[2], rgb[1], rgb[0])
    acc, base = K.hex_rgb(a.accent), K.hex_rgb(a.color)
    size = int(min(W, H) * 0.1 * a.size)
    marg = int(H * (1 - {"top": 0.2, "center": 0.5, "lower": 0.68, "bottom": 0.8}[a.position]) - size)
    lines = ["[Script Info]", "ScriptType: v4.00+", f"PlayResX: {W}", f"PlayResY: {H}", "WrapStyle: 0", "",
             "[V4+ Styles]",
             "Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, "
             "StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding",
             f"Style: Kenerate,{font_name},{size},{col(acc)},{col(base)},&H00000000,&H64000000,-1,0,0,0,100,100,0,0,1,"
             f"{max(2, size // 11)},0,2,40,40,{max(10, marg)},1", "", "[Events]",
             "Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text"]
    for p in pages:
        parts, t = [], p["start"]
        for w in p["words"]:
            d = max(1, int(round((w["end"] - t) * 100)))
            parts.append(f"{{\\k{d}}}{w['word'].upper() if a.style != 'minimal' and not a.keep_case else w['word']}")
            t = w["end"]
        lines.append(f"Dialogue: 0,{ts(p['start'])},{ts(p['end'])},Kenerate,,0,0,0,,{' '.join(parts)}")
    dst.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return dst


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--video", required=True, help="video URL or file")
    ap.add_argument("--words-json")
    ap.add_argument("--ass")
    ap.add_argument("--srt")
    ap.add_argument("--text", help="script text (timed evenly over the video: rough)")
    ap.add_argument("--offset", type=float, default=0.0, help="shift all captions by N seconds (+ later, - earlier)")
    ap.add_argument("--style", choices=STYLES, default="highlight")
    ap.add_argument("--max-words", type=int, default=3, help="words on screen at once (2-4 reads best on 9:16)")
    ap.add_argument("--max-chars", type=int, default=22)
    ap.add_argument("--position", choices=["top", "center", "lower", "bottom"], default="lower")
    ap.add_argument("--accent", default="#FFD60A")
    ap.add_argument("--color", default="#FFFFFF")
    ap.add_argument("--size", type=float, default=1.0, help="font size multiplier")
    ap.add_argument("--keep-case", action="store_true", help="don't uppercase")
    ap.add_argument("--fps", type=int, default=30)
    K.add_font_args(ap, "TTF/OTF for the captions")
    ap.add_argument("--preview-only", action="store_true", help="only write the contact sheet + .ass/.json (no video export)")
    ap.add_argument("--out", default="kenerate-out/captions-pro/captioned.mp4")
    ap.add_argument("--keep-temp", action="store_true")
    a = K.parse_args(ap)

    st = K.Status(4, "captions-pro")
    K.require(status=st)
    tmp = K.workdir(a.keep_temp)
    out = Path(a.out)
    out.parent.mkdir(parents=True, exist_ok=True)

    st.step("Fetching the video")
    src = K.fetch(a.video, tmp, "video", st)
    info = K.probe(src)
    W, H, D = info["width"], info["height"], info["duration"]
    if not (W and H and D):
        K.fail("Could not read the video size / length.", "Check the file plays.")
    W, H = W // 2 * 2, H // 2 * 2

    st.step("Word timing")
    if a.words_json:
        words, how = words_from_json(Path(a.words_json)), "exact (words json)"
    elif a.ass:
        words, how = words_from_ass(Path(a.ass)), "from ASS"
    elif a.srt:
        words, how = words_from_cues(K.parse_srt(Path(a.srt).read_text(encoding="utf-8", errors="replace"))), "estimated inside SRT cues"
    elif a.text:
        words, how = words_from_cues(K.auto_cues(a.text, D * 0.96, D * 0.02, 6)), "estimated over the whole video"
    else:
        K.fail("No caption source.", "Pass --words-json, --ass, --srt or --text.")
    for w in words:
        w["start"], w["end"] = round(w["start"] + a.offset, 3), round(w["end"] + a.offset, 3)
    words = [w for w in words if w["end"] > 0 and w["start"] < D]
    if not words:
        K.fail("No words inside the video's time range.", "Check --offset and that the captions match this video.")
    pages = pages_from_words(words, max(1, a.max_words), a.max_chars, 0.6)
    st.info(f"{len(words)} words in {len(pages)} caption pages; timing {how}")
    if "estimated" in how:
        st.warn("word times are estimated - for exact sync pass --words-json (e.g. ElevenLabs alignment or Whisper words)")

    st.step(f"Rendering '{a.style}' captions")
    r = Renderer(a, W, H)
    events = []
    for p in pages:
        if a.style == "typewriter":
            events.append({"start": p["start"], "end": p["end"], "anim_in": p["end"] - p["start"],
                           "draw": (lambda pp: lambda t: r.draw(pp, 0, t - pp["start"]))(p)})
            continue
        for k, w in enumerate(p["words"]):
            s0 = p["start"] if k == 0 else w["start"]
            e0 = p["words"][k + 1]["start"] if k + 1 < len(p["words"]) else p["end"]
            events.append({"start": s0, "end": e0, "anim_in": 0.22 if a.style in ("pop", "bounce", "highlight", "karaoke") else 0,
                           "draw": (lambda pp, kk, ss: lambda t: r.draw(pp, kk, t - ss))(p, k, w["start"])})
    lst = None if a.preview_only else K.timed_overlay_track(events, W, H, D, a.fps, tmp / "caps", st)
    stem = out.with_suffix("")
    font_name = Path(a.font).stem.split("-")[0] if a.font else "Arial"
    ass = write_ass(pages, stem.with_suffix(".ass"), W, H, a, font_name)
    wj = Path(str(stem) + "_words.json")
    wj.write_text(json.dumps([{k: v for k, v in w.items() if k != "cue_end"} for w in words], indent=1), encoding="utf-8")

    picks = pages if len(pages) <= 6 else [pages[int(i * (len(pages) - 1) / 5)] for i in range(6)]
    mids = [round(max(p["start"], min(p["words"][len(p["words"]) // 2]["start"] + 0.25, p["end"] - 0.06, D - 0.05)), 3)
            for p in picks] or [D / 2]  # a middle word: shows spoken vs upcoming words
    sheet_path = Path(str(stem) + "_preview.jpg")
    if a.preview_only:
        st.step("Preview only: compositing sample frames (no export)")
        cells = []
        for i, t in enumerate(mids):
            fr = K.open_image(K.extract_frame(src, t, tmp / f"pv_{i}.jpg")).convert("RGBA").resize((W, H))
            ev = next((e for e in events if e["start"] <= t < e["end"]), None)
            if ev:
                fr.alpha_composite(ev["draw"](t))
            p = tmp / f"pvc_{i}.jpg"
            fr.convert("RGB").save(p, quality=88)
            cells.append(p)
        sheet = K.contact_sheet(cells, sheet_path, cols=min(6, len(cells)), cell=300, labels=[f"{t:.1f}s" for t in mids])
        target = None
    else:
        st.step("Burning in")
        target = out
        K.overlay_track(src, lst, target, st)
        st.info("preview frames")
        sheet = K.video_preview(target, sheet_path, tmp=tmp / "pv", times=mids)
    K.finish({"output": K.file_info(target) if target else None, "preview": str(sheet.resolve()),
              "ass": str(ass.resolve()), "words_json": str(wj.resolve()), "style": a.style, "pages": len(pages),
              "words": len(words), "timing": how, "elapsed_s": st.elapsed(), "temp": str(tmp) if a.keep_temp else None})


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