#!/usr/bin/env python3
"""TRUE camera-shake detector (vidstab-style) via global optical flow. Our qcache `motion` is frame
absdiff — it conflates SUBJECT motion (bride walking = good) with CAMERA shake (handheld jitter =
bad). This separates them: estimate the DOMINANT global transform per frame (feature tracking +
RANSAC affine), which rejects moving-subject outliers and keeps only camera motion. Shake = JERK,
the frame-to-frame change in global motion: a steady pan has smooth motion (low jerk); a tremor has
erratic motion (high jerk). Cached per source like qcache.

  ./venv/bin/python shake_detect.py SOURCE_NAME      # one source (key in media_paths.json)
  ./venv/bin/python shake_detect.py --cam cam2       # all sources of a camera
  ./venv/bin/python shake_detect.py --all            # every source
Output: shakecache/{name}.json  {fps, frames:[{t, gmot, jerk}]}
"""
import json, subprocess, sys, os
import numpy as np
import cv2

FFMPEG = "/opt/homebrew/bin/ffmpeg"
W, FPS = 320, 10.0   # low res + sampled fps: shake is global, visible at 320px/10fps; ~cheap

FFPROBE = FFMPEG.replace("ffmpeg","ffprobe")
class DecodeError(Exception): pass

def _nframes_expected(path, fps=FPS):
    pr = subprocess.run([FFPROBE, "-v","error","-select_streams","v:0","-show_entries",
                         "format=duration","-of","json", path], capture_output=True, text=True)
    try: return int(float(json.loads(pr.stdout)["format"]["duration"]) * fps)
    except Exception: return None

def frames(path, w=W, fps=FPS, hw=True):
    """Yield grayscale frames (h,w) from ffmpeg at low res + sampled fps."""
    # probe height for the scaled width (keep aspect, even)
    pr = subprocess.run([FFPROBE, "-v","error","-select_streams","v:0",
                         "-show_entries","stream=width,height","-of","json", path],
                        capture_output=True, text=True)
    try:
        st = json.loads(pr.stdout)["streams"][0]; sw, sh = st["width"], st["height"]
    except Exception:
        sw, sh = 1920, 1080
    h = int(round(w * sh / sw)) // 2 * 2
    cmd = [FFMPEG, "-v","error"] + (["-hwaccel","videotoolbox"] if hw else []) + ["-i", path,
           "-vf", f"scale={w}:{h},fps={fps}", "-pix_fmt","gray","-f","rawvideo","-"]
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    sz = w * h; n = 0
    while True:
        buf = p.stdout.read(sz)
        if len(buf) < sz: break
        n += 1; yield np.frombuffer(buf, np.uint8).reshape(h, w)
    p.stdout.close(); err = (p.stderr.read() or b"").decode(errors="replace"); p.wait()
    # Só rejeita decode CURTO. VT ruidoso com frames completos NÃO força re-decode (ver analyze_visual).
    exp = _nframes_expected(path, fps)
    hard_fail = (
        n == 0
        or "No frame decoded" in err
        or "hardware accelerator failed" in err
        or (exp is not None and n < exp * 0.9)
        or (p.returncode != 0 and (exp is None or n < exp * 0.95))
    )
    if hard_fail:
        raise DecodeError(f"decodou {n}/{exp if exp else '?'} frames rc={p.returncode} hw={hw}")

def analyze(path):
    # Default software no batch (VIDEOTASK_SHAKE_HW=1 p/ VT).
    use_hw = os.environ.get("VIDEOTASK_SHAKE_HW", "").strip().lower() in ("1", "true", "yes")
    if use_hw:
        try:
            return _analyze_once(path, hw=True)
        except DecodeError as e:
            print(f"  ! decode hwaccel {os.path.basename(path)}: {e} — refazendo em software", flush=True)
            return _analyze_once(path, hw=False)
    return _analyze_once(path, hw=False)

def _analyze_once(path, hw=True):
    gm = []   # per-frame global translation magnitude (dominant shift = camera motion, FFT-robust to subject)
    prev = None
    han = None
    for f in frames(path, hw=hw):
        ff = f.astype(np.float32)
        if han is None or han.shape != ff.shape:
            han = cv2.createHanningWindow((ff.shape[1], ff.shape[0]), cv2.CV_32F)
        if prev is not None:
            (dx, dy), _ = cv2.phaseCorrelate(prev, ff, han)
            # clamp: a real inter-frame shift at 320px is at most ~30px; bigger = decode glitch
            # (corrupt NAL) — cap so one bad frame can't fake a huge jerk spike.
            gm.append(min(float(np.hypot(dx, dy)), 40.0))
        prev = ff
    gm = np.array(gm) if gm else np.array([0.0])
    # JERK = |change in global motion| = erratic vs smooth. Pan (smooth) -> low; tremor -> high.
    jerk = np.abs(np.diff(gm, prepend=gm[:1]))
    out = [{"t": round(i / FPS, 2), "gmot": round(float(gm[i]), 2), "jerk": round(float(jerk[i]), 2)}
           for i in range(len(gm))]
    return {"fps": FPS, "frames": out}

def sources_for(cam=None, name=None):
    import vlib as V
    M = json.load(open("media_paths.json"))
    P = json.load(open("placements.json"))["clips"]
    if name:
        p = M.get(name) or V.resolve_media_path(name, M)
        return {name: p} if p else {}
    res = {}
    for c in P:
        if cam is not None and c.get("cam") != cam:
            continue
        n = V.footage_key(c["name"], cam=c.get("cam"))
        path = V.resolve_media_path(c["name"], M, cam=c.get("cam"))
        if path and n:
            res[n] = path
    return res

if __name__ == "__main__":
    os.makedirs("shakecache", exist_ok=True)
    if sys.argv[1] == "--cam":
        srcs = sources_for(cam=sys.argv[2])
    elif sys.argv[1] == "--all":
        srcs = sources_for()
    else:
        srcs = sources_for(name=sys.argv[1])
    items = list(srcs.items())
    if "--rev" in sys.argv: items = items[::-1]   # reverso: worker extra converge com o forward
    failed = []
    for name, path in items:
        if not os.path.exists(path):
            print(f"  SKIP (no file): {name}"); continue
        out = f"shakecache/{name}.json"
        if os.path.exists(out):
            print(f"  cached: {name}", flush=True); continue
        try:
            r = analyze(path)
        except DecodeError as e:
            print(f"  ✗ DECODE FALHOU {name}: {e}", flush=True); failed.append(name); continue
        except Exception as e:
            print(f"  ERROR {name}: {e}", flush=True); failed.append(name); continue
        tmp = f"{out}.{os.getpid()}.tmp"
        json.dump(r, open(tmp, "w")); os.replace(tmp, out)   # atômico p/ workers concorrentes
        j = np.array([x["jerk"] for x in r["frames"]])
        g = np.array([x["gmot"] for x in r["frames"]])
        print(f"  {name:40} n={len(j):>5}  gmot med={np.median(g):5.1f}  "
              f"jerk med={np.median(j):4.1f} p90={np.percentile(j,90):5.1f} p99={np.percentile(j,99):6.1f}", flush=True)
    n_ok = len([1 for n, _ in items if os.path.exists(f"shakecache/{n}.json")])
    if n_ok == 0 and items:
        print(f"\n✗ SHAKE: 0/{len(items)} caches — decode total falhou: {failed}", flush=True)
        sys.exit(1)
    if failed:
        print(f"  ⚠️  SHAKE: {len(failed)} fonte(s) sem cache (cut segue parcial): {failed[:8]}",
              flush=True)
