#!/usr/bin/env python3
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# Crop a vertical video to 1:1, keeping a tracked object centered.
# usage: trackcrop.py [--no-audio] [--ak820 FPS] in.mp4 out.mp4 ; draw box around object on first frame, press Enter
# needs: opencv-contrib-python (not -headless), ffmpeg with h264_vaapi
import sys, subprocess, cv2, numpy as np

AK820_FRAMES = 255  # AJAZZ AK820 Pro screen holds at most 255 frames (see ak820-upload)
USAGE = "usage: trackcrop.py [--no-audio] [--ak820 FPS] in.mp4 out.mp4"

args = sys.argv[1:]
mute = "--no-audio" in args
if mute:
    args.remove("--no-audio")
kfps = None
if "--ak820" in args:
    i = args.index("--ak820")
    try:
        kfps = float(args[i + 1])
    except (IndexError, ValueError):
        sys.exit(USAGE)
    del args[i:i + 2]
if len(args) != 2:
    sys.exit(USAGE)
src, dst = args

def open_at(first):
    """Open src positioned at frame `first` (grab, not seek: exact on any codec)."""
    cap = cv2.VideoCapture(src)
    for _ in range(first):
        cap.grab()
    return cap

cap = cv2.VideoCapture(src)
fps = cap.get(cv2.CAP_PROP_FPS)
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # container estimate
cap.release()

# --ak820: keep a window that fits the keyboard's frame limit at FPS
first, count = 0, 10**9
if kfps and total / fps > AK820_FRAMES / kfps:
    dur, keep = total / fps, AK820_FRAMES / kfps
    ans = input(f"{dur:.1f}s clip, the keyboard holds {keep:.1f}s at {kfps:g} fps. "
                f"Keep [s]tart, [e]nd, or start at N seconds (0-{dur - keep:.1f}): ").strip().lower()
    try:
        t0 = 0.0 if ans in ("", "s") else dur - keep if ans == "e" else float(ans)
    except ValueError:
        sys.exit(f"not s, e or a number: {ans}")
    t0 = min(max(t0, 0.0), dur - keep)
    first, count = round(t0 * fps), int(keep * fps)
    print(f"keeping {t0:.2f}s to {t0 + keep:.2f}s")

cap = open_at(first)
ok, frame = cap.read()
if not ok:
    sys.exit(f"cannot read {src}")
h, w = frame.shape[:2]
s = min(w, h)
tr = cv2.TrackerCSRT_create()
cv2.namedWindow("select object", cv2.WINDOW_NORMAL)  # resizable, tall frames fit the screen
tr.init(frame, cv2.selectROI("select object", frame))
cv2.destroyAllWindows()

# pass 1: track centers (reuse last position when tracker loses it)
cx, cy = [], []
last = (w / 2, h / 2)
while ok and len(cx) < count:
    found, (x, y, bw, bh) = tr.update(frame)
    if found:
        last = (x + bw / 2, y + bh / 2)
    cx.append(last[0]); cy.append(last[1])
    ok, frame = cap.read()

# smooth path to kill jitter (~0.5s moving average)
k = max(int(fps // 2) | 1, 1)
smooth = lambda a: np.convolve(np.pad(a, k // 2, mode="edge"), np.ones(k) / k, "valid")
cx, cy = smooth(np.array(cx)), smooth(np.array(cy))

# pass 2: crop and pipe to ffmpeg, keep original audio
# ponytail: local ffmpeg has no libx264, encodes on GPU via VA-API; device path is hardcoded
ff = subprocess.Popen(["ffmpeg", "-y", "-vaapi_device", "/dev/dri/renderD128",
    "-f", "rawvideo", "-pix_fmt", "bgr24",
    "-s", f"{s}x{s}", "-r", str(fps), "-i", "-",
    *([] if mute else ["-ss", f"{first / fps}", "-t", f"{len(cx) / fps}", "-i", src,
                       "-map", "0:v", "-map", "1:a?", "-c:a", "copy"]),
    "-vf", "format=nv12,hwupload", "-c:v", "h264_vaapi", "-qp", "20", dst],
    stdin=subprocess.PIPE)
cap = open_at(first)
for x, y in zip(cx, cy):
    ok, frame = cap.read()
    if not ok:
        break
    x0 = int(np.clip(x - s / 2, 0, w - s))
    y0 = int(np.clip(y - s / 2, 0, h - s))
    ff.stdin.write(frame[y0:y0 + s, x0:x0 + s].tobytes())
ff.stdin.close()
ff.wait()