#!/usr/bin/env python3
"""Smart reframe: 16:9 -> 9:16 (or any aspect) that follows the action instead of a dumb centre crop.

How it works: black bars are removed first (cropdetect). Small grey frames are sampled a few times a second and
scored per column (or row) for motion (frame difference) and detail (edges), with a mild centre bias. The crop
window goes where the score is highest, shot by shot (hard cuts are detected and allowed to jump), smoothed with a
dead-zone so the camera doesn't wobble, then rendered with one ffmpeg crop expression. A preview sheet shows the
crop window on source frames before you trust it.

Examples
  python reframe.py --input https://.../wide.mp4 --aspect 9:16 --preview-only      # check the path first
  python reframe.py --input wide.mp4 --aspect 9:16 --style smooth
  python reframe.py --input wide.mp4 --aspect 9:16 --style static                  # one fixed framing per shot
  python reframe.py --input wide.mp4 --aspect 4:5 --keyframes "0:0.3,4.5:0.7"       # manual: time:position (0..1)
  python reframe.py --input wide.mp4 --aspect 9:16 --mode center
"""
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


def cropdetect(src: Path, D: float):
    """Black-bar crop (w, h, x, y) or None, from ffmpeg cropdetect on a few seconds spread over the clip."""
    exe = K.find_ffmpeg()
    import subprocess
    found = []
    for t in [D * f for f in (0.2, 0.5, 0.8)]:
        r = subprocess.run([exe, "-hide_banner", "-ss", f"{t:.2f}", "-i", str(src), "-t", "1", "-vf", "cropdetect=limit=24:round=2:reset=0",
                            "-an", "-f", "null", "-"], capture_output=True, text=True, errors="replace", timeout=120)
        m = re.findall(r"crop=(\d+):(\d+):(\d+):(\d+)", r.stderr)
        if m:
            found.append(tuple(int(v) for v in m[-1]))
    if not found:
        return None
    w = max(f[0] for f in found)
    h = max(f[1] for f in found)
    x = min(f[2] for f in found)
    y = min(f[3] for f in found)
    return (w, h, x, y)


def refine_cut(src: Path, t: float, fps: float, w: int, h: int) -> float:
    """Exact cut time: re-sample the gap before t at 30 fps and take the biggest jump."""
    lo = max(0.0, t - 1.0 / fps - 0.04)
    frames = list(K.gray_frames(src, 30, w, h, start=lo, dur=t - lo + 0.04))
    best, bt = -1.0, t
    for (t0, f0, _, _), (t1, f1, _, _) in zip(frames, frames[1:]):
        dd = K.frame_diff(f0, f1)
        if dd > best:
            best, bt = dd, t1
    return round(bt, 3)


def rdp(pts, eps):
    """Ramer-Douglas-Peucker simplification of [(t, x)]."""
    if len(pts) < 3:
        return pts
    (t0, x0), (t1, x1) = pts[0], pts[-1]
    dmax, idx = 0.0, 0
    for i in range(1, len(pts) - 1):
        t, x = pts[i]
        xl = x0 + (x1 - x0) * (t - t0) / max(t1 - t0, 1e-9)
        d = abs(x - xl)
        if d > dmax:
            dmax, idx = d, i
    if dmax > eps:
        return rdp(pts[:idx + 1], eps)[:-1] + rdp(pts[idx:], eps)
    return [pts[0], pts[-1]]


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--input", required=True)
    ap.add_argument("--aspect", default="9:16")
    ap.add_argument("--res", type=int, default=1080, help="short side of the output (default 1080)")
    ap.add_argument("--mode", choices=["track", "center", "fixed"], default="track")
    ap.add_argument("--focus", type=float, default=0.5, help="fixed mode: window position 0 (left / top) .. 1")
    ap.add_argument("--style", choices=["smooth", "static"], default="smooth",
                    help="smooth = gentle pans that follow the action; static = one framing per shot (no pans)")
    ap.add_argument("--keyframes", default="", help="manual path 't:pos,t:pos' (pos 0..1), overrides tracking")
    ap.add_argument("--smooth", type=float, default=0.8, help="0 (twitchy) .. 1 (very calm), default 0.8")
    ap.add_argument("--no-bars", action="store_true", help="skip black-bar detection")
    ap.add_argument("--analysis-fps", type=float, default=5.0)
    ap.add_argument("--preview-only", action="store_true", help="write the path + preview sheet, don't render")
    ap.add_argument("--out", default="kenerate-out/video-reframe/reframed.mp4")
    a = ap.parse_args()

    st = K.Status(5, "video-reframe")
    K.require(status=st)
    tmp = K.workdir()
    out = Path(a.out)
    out.parent.mkdir(parents=True, exist_ok=True)

    st.step("Fetching + probing")
    src = K.fetch(a.input, tmp, "input", st)
    info = K.probe(src)
    D, SW, SH = info["duration"], info["width"], info["height"]
    if not (D and SW and SH):
        K.fail("Could not read the video.", "Check the file plays.")
    bars = None if a.no_bars else cropdetect(src, D)
    if bars and (bars[0] < SW - 8 or bars[1] < SH - 8) and bars[0] > SW * 0.5 and bars[1] > SH * 0.5:
        st.info(f"black bars removed: {SW}x{SH} -> {bars[0]}x{bars[1]} at {bars[2]},{bars[3]}")
        CW, CH, OX, OY = bars
    else:
        bars, CW, CH, OX, OY = None, SW, SH, 0, 0
    TW, TH = K.size_for(a.aspect, a.res)
    tar = TW / TH
    if CW / CH > tar:  # source wider: move a window horizontally
        axis, win, span = "x", int(round(CH * tar / 2)) * 2, CW
        cw, ch = win, CH // 2 * 2
    else:
        axis, win, span = "y", int(round(CW / tar / 2)) * 2, CH
        cw, ch = CW // 2 * 2, win
    free = span - win
    st.info(f"{SW}x{SH} {D:.1f} s -> {TW}x{TH}; window {cw}x{ch}, moves along {axis} ({free}px of travel)")

    st.step("Analysing motion + detail" if a.mode == "track" and not a.keyframes else "Building the path")
    cuts, path = [], []
    if a.keyframes:
        for kv in a.keyframes.split(","):
            t, p = kv.split(":")
            path.append((float(t), K.clamp(float(p)) * free))
        path.sort()
        if path[0][0] > 0:
            path.insert(0, (0.0, path[0][1]))
    elif a.mode != "track" or free < 4:
        pos = (0.5 if a.mode == "center" else K.clamp(a.focus)) * free
        path = [(0.0, pos), (D, pos)]
    else:
        if axis == "x":
            aw, ah = 160, max(16, int(round(160 * CH / CW / 2)) * 2)
        else:
            aw, ah = max(16, int(round(160 * CW / CH / 2)) * 2), 160
        vf_src = src
        if bars:
            pre = tmp / "nobars.mp4"
            K.run_ffmpeg(["-i", str(src), "-vf", f"crop={CW}:{CH}:{OX}:{OY},scale={aw}:{ah}", "-an", "-c:v", "libx264" if " libx264 " in K.encoders() else "mpeg4",
                          "-q:v", "4", str(pre)], duration=D, status=st)
            vf_src = pre
        n_est = int(D * a.analysis_fps)
        prev, raw, last, diffs = None, [], None, []
        for i, (t, fr, w, h) in enumerate(K.gray_frames(vf_src, a.analysis_fps, aw, ah)):
            is_cut = False
            if prev is not None:
                dd = K.frame_diff(prev, fr)
                base = sorted(diffs[-10:])[len(diffs[-10:]) // 2] if diffs else 0.0
                if dd > 0.06 and dd > base * 3 + 0.02 and (last is None or t - last > 0.4):
                    cuts.append(refine_cut(vf_src, t, a.analysis_fps, aw, ah))
                    is_cut = True
                    last = t
                diffs.append(dd)
            prof = [0.0] * (w if axis == "x" else h)
            step = 2
            for y in range(0, h, step):
                row = y * w
                for x in range(0, w - 1, step):
                    v = fr[row + x]
                    e = abs(v - fr[row + x + 1]) + (abs(v - fr[row + x + w]) if y + 1 < h else 0)
                    m = abs(v - prev[row + x]) * 2.5 if prev is not None and not is_cut else 0
                    prof[x if axis == "x" else y] += e * 0.4 + m
            n = len(prof)
            cen = [prof[j] * (1.0 - 0.25 * abs(j / max(n - 1, 1) - 0.5) * 2) for j in range(n)]
            wn = max(1, int(round(n * win / span)))
            acc = sum(cen[:wn])
            best, bj = acc, 0
            for j in range(1, n - wn + 1):
                acc += cen[j + wn - 1] - cen[j - 1]
                if acc > best:
                    best, bj = acc, j
            target = bj / max(n - wn, 1) * free if n > wn else free / 2
            raw.append((t, target))
            prev = fr
            if n_est:
                st.progress(min(99, (i + 1) * 100 / n_est), f"t={t:.1f}s")
        if not raw:
            K.fail("No frames could be analysed.", "Check the video decodes (python kmedia.py probe FILE).")
        # split into shots, then smooth inside each shot
        bounds = [0.0] + cuts + [D + 1]
        dz = free * 0.06
        for si in range(len(bounds) - 1):
            seg = [(t, x) for t, x in raw if bounds[si] <= t < bounds[si + 1]]
            if not seg:
                continue
            if si > 0 and path:  # hold the old framing right up to the cut, then jump
                path.append((bounds[si] - 0.001, path[-1][1]))
            if a.style == "static":
                xs = sorted(x for _, x in seg)
                med = xs[len(xs) // 2]
                path += [(bounds[si], med), (seg[-1][0], med)]
                continue
            xs = [x for _, x in seg]
            med = [sorted(xs[max(0, i - 1):i + 2])[len(xs[max(0, i - 1):i + 2]) // 2] for i in range(len(xs))]  # drop spikes
            k = max(1, int(round(a.analysis_fps * (0.2 + 1.3 * K.clamp(a.smooth)))))  # centred window: no lag
            n_ = len(med)
            sm = []
            for i in range(n_):  # symmetric window that shrinks at shot edges (no pull towards older positions)
                kk = min(k, i, n_ - 1 - i)
                sm.append(sum(med[i - kk:i + kk + 1]) / (2 * kk + 1))
            held, cur = [], sm[0]
            for v in sm:  # dead-zone: ignore drifts smaller than ~6% of the travel
                if abs(v - cur) > dz * 0.5:
                    cur = v
                held.append(cur)
            sm2 = [(seg[i][0], held[i]) for i in range(len(seg))]
            sm2[0] = (bounds[si], sm2[0][1])
            path += sm2
        if path[0][0] > 0:
            path.insert(0, (0.0, path[0][1]))
        path.append((D, path[-1][1]))
    path = [(round(t, 3), round(min(max(x, 0), free), 1)) for t, x in path]
    # simplify, keep the expression short
    eps = max(1.0, free * 0.004)
    simp = rdp(path, eps)
    while len(simp) > 120:
        eps *= 1.6
        simp = rdp(path, eps)
    pts = simp
    st.info(f"path: {len(path)} samples -> {len(pts)} keyframes; {len(cuts)} cut(s) at {', '.join(f'{c:.1f}s' for c in cuts[:8]) or 'none'}")

    st.step("Preview (crop window on source frames)")
    from PIL import ImageDraw

    def pos_at(t):
        if t <= pts[0][0]:
            return pts[0][1]
        for (t0, x0), (t1, x1) in zip(pts, pts[1:]):
            if t0 <= t <= t1:
                return x0 + (x1 - x0) * (t - t0) / max(t1 - t0, 1e-9)
        return pts[-1][1]

    cells, labels = [], []
    for i, t in enumerate([D * f for f in (0.05, 0.25, 0.45, 0.65, 0.85, 0.97)]):
        fr = K.open_image(K.extract_frame(src, t, tmp / f"pv_{i}.jpg")).convert("RGB")
        d = ImageDraw.Draw(fr, "RGBA")
        p = pos_at(t)
        x0, y0 = (OX + p, OY) if axis == "x" else (OX, OY + p)
        x1, y1 = x0 + cw, y0 + ch
        d.rectangle((0, 0, fr.width, fr.height), outline=None)
        for box in ((0, 0, x0, fr.height), (x1, 0, fr.width, fr.height)) if axis == "x" else ((0, 0, fr.width, y0), (0, y1, fr.width, fr.height)):
            if box[2] > box[0] and box[3] > box[1]:
                d.rectangle(box, fill=(0, 0, 0, 150))
        d.rectangle((x0, y0, x1, y1), outline=(255, 214, 10, 255), width=max(3, SW // 300))
        pth = tmp / f"pvb_{i}.jpg"
        fr.save(pth, quality=85)
        cells.append(pth)
        labels.append(f"{t:.1f}s")
    sheet = K.contact_sheet(cells, out.with_name(out.stem + "_path_preview.jpg"), cols=3, cell=420, labels=labels)
    pj = out.with_name(out.stem + "_path.json")
    pj.write_text(json.dumps({"axis": axis, "window": [cw, ch], "travel_px": free, "bars": bars, "cuts": cuts,
                              "keyframes": [{"t": t, "pos": round(x / free, 4) if free else 0.5} for t, x in pts]}, indent=1), encoding="utf-8")
    if a.preview_only:
        K.finish({"preview": str(sheet.resolve()), "path_json": str(pj.resolve()), "rendered": False,
                  "hint": "Looks right? Re-run without --preview-only. Wrong subject? Pass --keyframes from path_json (edit pos).",
                  "elapsed_s": st.elapsed()})
        return

    st.step(f"Rendering {TW}x{TH}")
    terms = [f"{pts[0][1]:.1f}"]
    for (t0, x0), (t1, x1) in zip(pts, pts[1:]):
        dx = x1 - x0
        if abs(dx) < 0.05:
            continue
        terms.append(f"{dx:+.1f}*clip((t-{t0:.3f})/{max(t1 - t0, 0.001):.3f},0,1)")
    expr = "".join(terms)
    pos_expr = f"min(max({expr},0),{free})"
    cx, cy = (f"{OX}+{pos_expr}", str(OY)) if axis == "x" else (str(OX), f"{OY}+{pos_expr}")
    vf = f"crop={cw}:{ch}:'{cx}':'{cy}',scale={TW}:{TH}:flags=lanczos,setsar=1,format=yuv420p"
    K.run_ffmpeg(["-i", str(src), "-vf", vf, "-map", "0:v:0", "-map", "0:a?"] + K.video_codec_args(19) + ["-c:a", "copy", str(out)],
                 duration=D, status=st)
    K.finish({"output": K.file_info(out), "preview": str(sheet.resolve()), "path_json": str(pj.resolve()), "cuts": cuts,
              "keyframes": len(pts), "mode": a.mode if not a.keyframes else "keyframes", "elapsed_s": st.elapsed()})


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