#!/usr/bin/env python3
"""Ad recreator: rebuild a reference ad's STRUCTURE (shot count, shot lengths, pacing, framing, colour, where the text
sits) with the user's own product - never its logos, people, music or copy.

  analyze   reference ad -> shot breakdown (cuts, durations, keyframes, colour, brightness, motion, text bands,
            a role guess per shot: hook / problem / demo / proof / cta), a numbered contact sheet and a plan.json
            template the agent fills in with new prompts and, after generating, the new shot urls
  assemble  plan.json with "src" per shot -> the new ad cut to the reference's timing, optional on-screen text per
            shot, optional music / voice-over (only ones the user asked for), and a side-by-side check sheet

Examples
  python recreate.py analyze --ref https://.../reference_ad.mp4 --name coldbrew
  python recreate.py assemble --plan kenerate-out/ad-recreator/coldbrew_plan.json --aspect 9:16
"""
from __future__ import annotations

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

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


def loudness(src: Path) -> float:
    r = subprocess.run([K.find_ffmpeg(), "-hide_banner", "-i", str(src), "-af", "volumedetect", "-vn", "-f", "null", "-"],
                       capture_output=True, text=True, errors="replace", timeout=300)
    m = re.search(r"mean_volume:\s*(-?[\d.]+) dB", r.stderr)
    return float(m.group(1)) if m else -99.0


def role_guess(i: int, n: int, start: float, D: float) -> str:
    if i == 0 or start < 2.5:
        return "hook"
    if i == n - 1 or start > D * 0.85:
        return "cta"
    f = start / D
    return "problem" if f < 0.35 else ("demo" if f < 0.7 else "proof")


def analyze(a) -> None:
    st = K.Status(5, "ad-recreator analyze")
    K.require(status=st)
    tmp = K.workdir()
    out_dir = Path(a.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    st.step("Fetching the reference ad")
    src = K.fetch(a.ref, tmp, "reference", st)
    info = K.probe(src)
    D = info["duration"]
    if D <= 0.5 or not info["has_video"]:
        K.fail("The reference is not a readable video.", "Pass the ad as an MP4 / MOV link or file.")
    st.info(f"{info['width']}x{info['height']}, {D:.1f} s, audio: {'yes' if info['has_audio'] else 'no'}")

    st.step("Finding the cuts")
    cuts, _ = K.scene_cuts(src, threshold=a.threshold, fps=8.0)
    bounds = [0.0] + [c for c in cuts if 0.3 < c < D - 0.3] + [D]
    shots = [(bounds[i], bounds[i + 1]) for i in range(len(bounds) - 1) if bounds[i + 1] - bounds[i] > 0.2]
    st.info(f"{len(shots)} shot(s), average {D / max(len(shots), 1):.2f} s per shot")

    st.step(f"Measuring {len(shots)} shot(s)")
    from PIL import ImageStat
    rows, frames = [], []
    for i, (s0, s1) in enumerate(shots):
        mid = s0 + (s1 - s0) * 0.5
        kf = K.extract_frame(src, mid, out_dir / f"{a.name}_ref_shot{i + 1:02d}.jpg")
        frames.append(kf)
        im = K.open_image(kf).convert("RGB")
        W, H = im.size
        mean = tuple(int(v) for v in ImageStat.Stat(im).mean[:3])
        bright = round(sum(mean) / 765, 3)
        bands = {}
        for name, box in (("top", (0, 0, W, int(H * 0.25))), ("middle", (0, int(H * 0.375), W, int(H * 0.625))),
                          ("bottom", (0, int(H * 0.75), W, H))):
            _, busy = K.region_stats(im, box)
            bands[name] = busy
        text_band = max(bands, key=bands.get) if max(bands.values()) > 0.35 else None
        motion = 0.0
        fr = list(K.gray_frames(src, 6, 64, start=s0, dur=min(s1 - s0, 2.0)))
        if len(fr) > 1:
            motion = round(sum(K.frame_diff(x[1], y[1]) for x, y in zip(fr, fr[1:])) / (len(fr) - 1), 4)
        rows.append({"shot": i + 1, "start": round(s0, 2), "end": round(s1, 2), "seconds": round(s1 - s0, 2),
                     "role": role_guess(i, len(shots), s0, D), "avg_color": K.rgb_hex(mean), "brightness": bright,
                     "motion": "static" if motion < 0.01 else ("moving" if motion < 0.05 else "fast"),
                     "busy_band": text_band, "keyframe": str(kf.resolve())})
        st.progress((i + 1) * 100 / len(shots), f"shot {i + 1}/{len(shots)}")

    st.step("Contact sheet + plan template")
    sheet = K.contact_sheet(frames, out_dir / f"{a.name}_ref_sheet.jpg", cols=min(6, len(frames)), cell=260,
                            labels=[f"{r['shot']}. {r['role']} {r['seconds']}s" for r in rows])
    vol = loudness(src) if info["has_audio"] else None
    aspect = "9:16" if info["height"] > info["width"] * 1.2 else ("1:1" if abs(info["height"] - info["width"]) < info["width"] * 0.1 else "16:9")
    breakdown = {"reference": a.ref, "duration": round(D, 2), "size": f"{info['width']}x{info['height']}", "aspect": aspect,
                 "shots": rows, "cuts_per_10s": round(len(shots) / max(D, 0.1) * 10, 2),
                 "audio": {"present": info["has_audio"], "mean_volume_db": vol}}
    bj = out_dir / f"{a.name}_breakdown.json"
    bj.write_text(json.dumps(breakdown, indent=1), encoding="utf-8")
    plan = {"name": a.name, "aspect": aspect, "reference": a.ref,
            "audio": {"music": "", "voiceover": "", "note": "fill only if the user asked for music / a voice-over; never reuse the reference's audio"},
            "shots": [{"shot": r["shot"], "seconds": r["seconds"], "role": r["role"], "reference_look": f"{r['motion']}, {r['avg_color']}, text in {r['busy_band'] or 'none'}",
                       "ref_keyframe": r["keyframe"], "prompt": "", "text": "", "src": ""} for r in rows]}
    pj = out_dir / f"{a.name}_plan.json"
    pj.write_text(json.dumps(plan, indent=1), encoding="utf-8")
    st.step("Done")
    K.finish({"breakdown": str(bj.resolve()), "plan": str(pj.resolve()), "sheet": str(sheet.resolve()), "shots": len(rows),
              "duration": round(D, 2), "aspect": aspect, "elapsed_s": st.elapsed()})


def assemble(a) -> None:
    plan = json.loads(Path(a.plan).read_text(encoding="utf-8"))
    shots = plan.get("shots", [])
    missing = [s["shot"] for s in shots if not s.get("src")]
    if missing:
        K.fail(f"Shots without a 'src': {missing}", "Generate those shots and put their Kenerate urls in the plan, or remove them.")
    st = K.Status(6, "ad-recreator assemble")
    K.require(status=st)
    tmp = K.workdir(a.keep_temp)
    name = plan.get("name", "ad")
    out = Path(a.out or f"kenerate-out/ad-recreator/{name}_recreated.mp4")
    out.parent.mkdir(parents=True, exist_ok=True)
    W, H = K.size_for(a.aspect or plan.get("aspect", "9:16"), a.res)

    st.step(f"Fetching {len(shots)} shot(s)")
    srcs = K.fetch_all([s["src"] for s in shots], tmp / "in", st, prefix="shot")
    st.step(f"Cutting each shot to the reference timing ({W}x{H})")
    norm = []
    for i, (p, s) in enumerate(zip(srcs, shots), 1):
        dst = tmp / f"n{i:02d}.mp4"
        want = float(s.get("seconds") or 2.0)
        got = K.normalize_clip(p, dst, W, H, a.fps, float(s.get("in", 0.0)), want, "cover", mute=not a.keep_shot_audio,
                               still_seconds=want, move=i)
        if got < want - 0.1:
            st.warn(f"shot {i} is {got:.2f}s, shorter than the reference's {want:.2f}s: the ad gets shorter")
        norm.append((dst, got))
        st.progress(i * 100 / len(srcs), f"shot {i}/{len(srcs)}")
    st.step("Joining (hard cuts, like the reference)")
    joined = tmp / "joined.mp4"
    total = K.join_clips(norm, joined, a.transition, a.xfade, st)
    cur = joined
    st.step("On-screen text")
    cues, t = [], 0.0
    for (p, d), s in zip(norm, shots):
        if s.get("text"):
            cues.append((t + 0.1, t + d - 0.05, s["text"]))
        t += d - (0 if a.transition == "cut" else a.xfade)
    if cues:
        lst = K.caption_track(cues, W, H, total, tmp / "caps", a.text_style, a.text_pos, a.font, status=st)
        capped = tmp / "capped.mp4"
        K.overlay_track(cur, lst, capped, st)
        cur = capped
    st.step("Audio")
    aud = plan.get("audio", {})
    music, vo = a.music or aud.get("music"), a.voiceover or aud.get("voiceover")
    if vo:
        vp = K.fetch(vo, tmp / "in", "voiceover", st)
        voiced = tmp / "voiced.mp4"
        K.run_ffmpeg(["-i", str(cur), "-i", str(vp), "-filter_complex",
                      "[0:a]volume=0.25[c];[1:a]aresample=48000,aformat=channel_layouts=stereo[v];"
                      "[c][v]amix=inputs=2:duration=first:normalize=0,alimiter=limit=0.95[a]",
                      "-map", "0:v", "-map", "[a]", "-c:v", "copy"] + K.AUDIO_ARGS + [str(voiced)], duration=total, status=st)
        cur = voiced
    if music:
        mp = K.fetch(music, tmp / "in", "music", st)
        mixed = tmp / "mixed.mp4"
        K.mix_music(cur, mp, mixed, a.music_db, duck=True, keep_voice=True, status=st)
        cur = mixed
    if not music and not vo:
        st.info("no music / voice-over (as agreed with the user)" if not a.keep_shot_audio else "keeping the shots' own sound")
    K.run_ffmpeg(["-i", str(cur), "-c", "copy", "-movflags", "+faststart", str(out)])

    st.step("Check sheet (reference vs new, shot by shot)")
    cells, labels, t = [], [], 0.0
    for i, ((p, d), s) in enumerate(zip(norm, shots), 1):
        ref = Path(s["ref_keyframe"]) if s.get("ref_keyframe") and Path(s["ref_keyframe"]).is_file() else None
        if ref is None:
            ref = next(iter(sorted(out.parent.glob(f"{name}_ref_shot{i:02d}.jpg"))), None)
        new = K.extract_frame(out, t + d / 2, tmp / f"new_{i:02d}.jpg")
        if ref:
            cells.append(ref)
            labels.append(f"ref {i} {s.get('role', '')}")
        cells.append(new)
        labels.append(f"new {i} {d:.1f}s")
        t += d
    sheet = K.contact_sheet(cells, out.with_name(out.stem + "_check.jpg"), cols=6 if len(cells) > 6 else len(cells), cell=240, labels=labels)
    K.finish({"output": K.file_info(out), "check_sheet": str(sheet.resolve()), "seconds": round(total, 2), "shots": len(norm),
              "music": bool(music), "voiceover": bool(vo), "elapsed_s": st.elapsed()})


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    an = sub.add_parser("analyze", help="break a reference ad into shots + a plan template")
    an.add_argument("--ref", required=True, help="reference ad (URL or file) the user may use for reference")
    an.add_argument("--name", default="ad")
    an.add_argument("--threshold", type=float, default=0.06, help="cut sensitivity 0..1 (lower = more cuts)")
    an.add_argument("--out-dir", default="kenerate-out/ad-recreator")
    asb = sub.add_parser("assemble", help="cut the new shots to the reference's timing")
    asb.add_argument("--plan", required=True)
    asb.add_argument("--aspect", default="")
    asb.add_argument("--res", type=int, default=1080)
    asb.add_argument("--fps", type=int, default=30)
    asb.add_argument("--transition", default="cut")
    asb.add_argument("--xfade", type=float, default=0.25)
    asb.add_argument("--text-style", choices=K.CAPTION_STYLES, default="box")
    asb.add_argument("--text-pos", choices=["top", "center", "bottom"], default="top")
    asb.add_argument("--music", default="", help="music the user asked for (url / file)")
    asb.add_argument("--music-db", type=float, default=-14.0)
    asb.add_argument("--voiceover", default="", help="voice-over the user asked for (url / file)")
    asb.add_argument("--keep-shot-audio", action="store_true", help="keep the generated shots' native audio")
    K.add_font_args(asb, "TTF/OTF for on-screen text")
    asb.add_argument("--out", default="")
    asb.add_argument("--keep-temp", action="store_true")
    a = K.parse_args(ap)
    analyze(a) if a.cmd == "analyze" else assemble(a)


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