#!/usr/bin/env python3
"""B-roll assembler: lay generated b-roll cutaways over an A-roll (talking head, avatar, podcast video) or over a
voice-over track, keeping the A-roll's sound. B-roll can be timed by seconds, by a keyword in the SRT, or spread
automatically into the gaps; full-frame cutaways or picture-in-picture.

Plan file (recommended, validate first with --dry-run):
  {"aroll": "talk.mp4", "srt": "talk.srt",
   "broll": [{"src": "https://.../beans.mp4", "keyword": "beans", "dur": 2.5},
             {"src": "https://.../pour.mp4", "at": 7.2, "dur": 3, "in": 1.0},
             {"src": "https://.../cafe.png"}]}

Examples
  python broll.py --plan plan.json --dry-run                      # check timings, no render
  python broll.py --aroll talk.mp4 --broll "beans.mp4::at=2,dur=2.5" --broll "cup.png::keyword=morning" --srt talk.srt
  python broll.py --aroll voiceover.mp3 --broll a.mp4 --broll b.mp4 --broll c.png --aspect 9:16   # VO-only edit
"""
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


def parse_item(s: str) -> dict:
    src, _, opts = s.partition("::")
    it: dict = {"src": src.strip()}
    for kv in [x for x in opts.split(",") if x.strip()]:
        k, _, v = kv.partition("=")
        k = k.strip()
        it[k] = v.strip() if k in ("keyword", "mode") else float(v)
    return it


def schedule(items, D: float, cues, min_gap: float, default_dur: float, skip_start: float, vo_only: bool):
    """Give every item a start time and a duration. Returns (items, warnings)."""
    warns = []
    for it in items:
        if "keyword" in it and "at" not in it:
            kw = str(it["keyword"]).lower()
            hit = next((c for c in cues if kw in c[2].lower()), None)
            if hit:
                it["at"] = round(hit[0], 3)
                it.setdefault("dur", round(min(max(hit[1] - hit[0], 1.5), 4.0), 3))
            else:
                warns.append(f"keyword '{it['keyword']}' not found in the captions: placed automatically")
    if vo_only:  # back to back over the whole voice-over
        free = [it for it in items if "at" not in it]
        fixed = sum(it.get("dur", 0) for it in free if "dur" in it)
        n_auto = sum(1 for it in free if "dur" not in it)
        rest = max(D - fixed, 0)
        t = 0.0
        for it in free:
            it.setdefault("dur", round(rest / max(n_auto, 1), 3))
            it["at"] = round(t, 3)
            t += it["dur"]
    else:
        auto = [it for it in items if "at" not in it]
        if auto:  # free gaps between the fixed shots, largest first, then filled in time order
            gaps = [(min(skip_start, D * 0.3), D - 0.8)]
            for it in items:
                if "at" in it:
                    a0, a1 = it["at"] - min_gap, it["at"] + it.get("dur", default_dur) + min_gap
                    gaps = [g for x in gaps for g in ((x[0], min(x[1], a0)), (max(x[0], a1), x[1])) if g[1] - g[0] > 0.05]
            slots = []
            for _ in auto:
                if not gaps:
                    break
                g = max(gaps, key=lambda x: x[1] - x[0])
                gaps.remove(g)
                mid = (g[0] + g[1]) / 2
                slots.append(g)
                gaps += [(g[0], mid - 0.001), (mid + 0.001, g[1])] if g[1] - g[0] > 2 * default_dur else []
            slots.sort()
            if len(slots) < len(auto):
                warns.append(f"only room for {len(slots)} of {len(auto)} auto-placed b-roll shots: the rest were dropped")
            for it, g in zip(auto, slots):
                it.setdefault("dur", min(default_dur, (g[1] - g[0]) * 0.8))
                it["at"] = round(g[0] + max(0.0, (g[1] - g[0] - it["dur"]) / 2), 3)
            items[:] = [it for it in items if "at" in it]
    for it in items:
        it.setdefault("dur", default_dur)
        it["dur"] = round(max(0.5, min(it["dur"], D - it["at"])), 3)
    items.sort(key=lambda x: x["at"])
    for a, b in zip(items, items[1:]):  # no overlaps: trim the earlier one
        if a["at"] + a["dur"] > b["at"] - (0 if vo_only else min_gap):
            new = round(b["at"] - (0 if vo_only else min_gap) - a["at"], 3)
            if new < 0.5:
                warns.append(f"b-roll at {b['at']:.2f}s overlaps the previous one: it was shortened a lot ({new:.2f}s)")
            a["dur"] = max(0.4, new)
    if vo_only:
        covered = sum(it["dur"] for it in items)
        if covered < D - 0.2:
            warns.append(f"b-roll covers {covered:.1f}s of {D:.1f}s: the last shot is held to the end")
            items[-1]["dur"] = round(D - items[-1]["at"], 3)
    return items, warns


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--plan", help="JSON plan (see above)")
    ap.add_argument("--aroll", help="A-roll video, or a voice-over audio file for a b-roll-only edit")
    ap.add_argument("--broll", action="append", default=[], help="SRC or 'SRC::at=3.2,dur=2.5,in=1,keyword=coffee,mode=pip'")
    ap.add_argument("--srt", help="captions of the A-roll, for keyword placement")
    ap.add_argument("--aspect", default="", help="output aspect (default: the A-roll's; 9:16 for a voice-over)")
    ap.add_argument("--res", type=int, default=1080)
    ap.add_argument("--fps", type=int, default=30)
    ap.add_argument("--mode", choices=["cutaway", "pip"], default="cutaway", help="full-frame cutaways or picture-in-picture")
    ap.add_argument("--pip-pos", choices=["tl", "tr", "bl", "br"], default="tr")
    ap.add_argument("--fade", type=float, default=0.15, help="b-roll fade in / out seconds")
    ap.add_argument("--default-dur", type=float, default=2.5)
    ap.add_argument("--skip-start", type=float, default=2.0, help="keep the first N s on the A-roll (the hook)")
    ap.add_argument("--dry-run", action="store_true", help="print the timeline and stop")
    ap.add_argument("--out", default="kenerate-out/b-roll-assembler/broll_edit.mp4")
    ap.add_argument("--keep-temp", action="store_true")
    a = ap.parse_args()

    plan = json.loads(Path(a.plan).read_text(encoding="utf-8")) if a.plan else {}
    aroll = a.aroll or plan.get("aroll")
    items = [dict(x) for x in plan.get("broll", [])] + [parse_item(s) for s in a.broll]
    srt = a.srt or plan.get("srt")
    if not aroll or not items:
        K.fail("Need an A-roll (or voice-over) and at least one b-roll.", "Pass --aroll and --broll (repeat), or --plan plan.json.")

    st = K.Status(5, "b-roll-assembler")
    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 A-roll")
    ar = K.fetch(aroll, tmp / "in", "aroll", st)
    ainfo = K.probe(ar)
    D = ainfo["duration"]
    vo_only = not ainfo["has_video"] or K.kind_of(str(ar)) == "audio"
    aspect = a.aspect or plan.get("aspect") or ("9:16" if vo_only else f"{ainfo['width']}x{ainfo['height']}")
    if not a.aspect and not vo_only:
        W, H = K.size_for(aspect)
        s = a.res / min(W, H)
        W, H = int(W * s) // 2 * 2, int(H * s) // 2 * 2
    else:
        W, H = K.size_for(aspect, a.res)
    cues = K.parse_srt(Path(srt).read_text(encoding="utf-8", errors="replace")) if srt else []
    items, warns = schedule(items, D, cues, 0.3, a.default_dur, a.skip_start, vo_only)
    for w in warns:
        st.warn(w)
    st.info(f"{'voice-over' if vo_only else 'A-roll'} {D:.1f} s -> {W}x{H}; {len(items)} b-roll shot(s):")
    for i, it in enumerate(items, 1):
        st.info(f"  {i:2d}. {it['at']:6.2f}s +{it['dur']:.2f}s  {it.get('mode', a.mode):7s} {str(it['src'])[-60:]}")
    if a.dry_run:
        K.finish({"dry_run": True, "duration": round(D, 2), "frame": f"{W}x{H}", "timeline": items, "warnings": warns})
        return

    st.step(f"Fetching + normalizing {len(items)} b-roll shot(s)")
    for i, it in enumerate(items, 1):
        p = K.fetch(str(it["src"]), tmp / "in", f"b{i:02d}", st)
        mode = it.get("mode", a.mode)
        bw, bh = (W, H) if mode == "cutaway" else (int(W * 0.42) // 2 * 2, int(H * 0.42) // 2 * 2)
        dst = tmp / f"b{i:02d}.mp4"
        K.normalize_clip(p, dst, bw, bh, a.fps, float(it.get("in", 0.0)), it["dur"], "cover", mute=True,
                         still_seconds=it["dur"], move=i)
        it.update({"_path": dst, "_w": bw, "_h": bh, "_mode": mode})
        st.progress(i * 100 / len(items), f"{i}/{len(items)}")

    st.step("Base track")
    base = tmp / "base.mp4"
    if vo_only:
        K.run_ffmpeg(["-f", "lavfi", "-i", f"color=c=black:s={W}x{H}:r={a.fps}:d={D:.3f}", "-i", str(ar),
                      "-map", "0:v", "-map", "1:a"] + K.video_codec_args(18, "veryfast") + K.AUDIO_ARGS + ["-t", f"{D:.3f}", str(base)],
                     duration=D, status=st)
    else:
        K.normalize_clip(ar, base, W, H, a.fps, status=st)

    st.step("Compositing cutaways")
    inputs, parts, last = ["-i", str(base)], [], "[0:v]"
    m = int(min(W, H) * 0.04)
    for i, it in enumerate(items, 1):
        inputs += ["-i", str(it["_path"])]
        f = min(a.fade, it["dur"] / 3)
        fx = (f"fade=in:st=0:d={f:.3f}:alpha=1,fade=out:st={max(it['dur'] - f, 0):.3f}:d={f:.3f}:alpha=1," if f > 0 and not vo_only else "")
        if it["_mode"] == "pip":
            border = max(4, W // 200)
            parts.append(f"[{i}:v]tpad=stop_mode=clone:stop_duration={it['dur']:.3f},pad=iw+{border * 2}:ih+{border * 2}:{border}:{border}:white,format=yuva420p,{fx}"
                         f"setpts=PTS-STARTPTS+{it['at']:.3f}/TB[b{i}]")
            px = {"tl": f"{m}", "bl": f"{m}", "tr": f"W-w-{m}", "br": f"W-w-{m}"}[a.pip_pos]
            py = {"tl": f"{m}", "tr": f"{m}", "bl": f"H-h-{m}", "br": f"H-h-{m}"}[a.pip_pos]
        else:
            parts.append(f"[{i}:v]tpad=stop_mode=clone:stop_duration={it['dur']:.3f},format=yuva420p,{fx}setpts=PTS-STARTPTS+{it['at']:.3f}/TB[b{i}]")
            px, py = "0", "0"
        parts.append(f"{last}[b{i}]overlay={px}:{py}:eof_action=pass:enable='between(t,{it['at']:.3f},{it['at'] + it['dur']:.3f})'[v{i}]")
        last = f"[v{i}]"
    graph = ";".join(parts) + f";{last}format=yuv420p[v]"
    K.run_ffmpeg(inputs + ["-filter_complex", graph, "-map", "[v]", "-map", "0:a"] + K.video_codec_args(19) + ["-c:a", "copy", str(out)],
                 duration=D, status=st)

    st.step("Preview + timeline")
    times = [it["at"] + it["dur"] / 2 for it in items][:5]
    times = sorted(set([max(0.1, min(a.skip_start / 2, D * 0.05))] + times))
    sheet = K.video_preview(out, out.with_name(out.stem + "_preview.jpg"), tmp=tmp / "pv", times=times)
    csv_path = out.with_name(out.stem + "_timeline.csv")
    with open(csv_path, "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["#", "start", "end", "mode", "source"])
        for i, it in enumerate(items, 1):
            w.writerow([i, f"{it['at']:.2f}", f"{it['at'] + it['dur']:.2f}", it["_mode"], it["src"]])
    K.finish({"output": K.file_info(out), "preview": str(sheet.resolve()), "timeline_csv": str(csv_path.resolve()),
              "shots": [{k: v for k, v in it.items() if not k.startswith("_")} for it in items], "warnings": warns,
              "elapsed_s": st.elapsed()})


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