#!/usr/bin/env python3
"""Find the best face in a video around a timestamp and write a square crop.

These clips are mostly SCREEN SHARES -- Jeremy's webcam is a small inset in one
corner, and the middle of the frame is a browser window. A centre crop therefore
returns a screenshot of a news article, which is what the first thumbnail did.
So the face has to be located, not assumed.

  facecrop.py <video> <time> <out.jpg> <size> [window] [step] [aspect]
  facecrop.py --detect <video> <time> [window] [step] [aspect]
  facecrop.py --crop x,y,w,h <video> <time> <out.jpg> <size> [aspect]

Samples a few frames around <time>, runs YuNet on each, and keeps the largest,
most confident face. Prints "x0 y0 side score t cwid ch" or "NONE" -- the first
five fields are the original contract and must not move; cwid/ch were appended
so a caller can RECORD the box that was cut instead of inferring it.

--detect runs the detection and prints the face box, the auto crop and the frame
size as JSON. It WRITES NOTHING. This is what the face judger asks, so a person
can see where the detector looked before deciding whether the crop it implies is
the picture they want.

--crop x,y,w,h skips detection entirely and cuts exactly that box at exactly
<time>. This is how a human's framing survives into a rebuild: the clamp below
preserves the crop's SIZE at the cost of moving it OFF the face, which is how a
corner ends up showing YouTube chrome, and re-deriving a framing that was chosen
to avoid that would silently undo the choice. A hand-framed crop reports a
confidence of -1.000 -- it was chosen, not detected, and reporting 1.000 would
be claiming a certainty no detector produced.

<aspect> is width/height, default 1.0 (square). Pass 1.7778 for 16:9. The output
is <size> wide and <size>/<aspect> tall, so `480 ... 1.7778` fills a 480x270
secondary box in the shorts layout exactly. The crop is taken about the same face
centre either way -- only the box around it changes shape -- so the two options
are a fair comparison rather than two different framings.
"""
import sys, os, json, subprocess, tempfile
import cv2
import numpy as np

MODEL = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                     "..", "models", "yunet.onnx")
MODEL = os.path.normpath(MODEL)

def frames(video, t0, window, step):
    ts = []
    t = t0 - window
    while t <= t0 + window + 1e-9:
        if t >= 0:
            ts.append(round(t, 3))
        t += step
    return ts

def grab(video, t, path):
    subprocess.run(["ffmpeg", "-nostdin", "-v", "error", "-y", "-ss", str(t),
                    "-i", video, "-frames:v", "1", path],
                   check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return os.path.exists(path) and os.path.getsize(path) > 0

def auto_crop(x, y, fw, fh, w, h, aspect):
    """The crop facecrop.py takes around a face box, and its geometry alone.

    Lifted out of main() unchanged so --detect can report the SAME box the
    normal path cuts without running the cut. lib/face-types.ts re-implements
    exactly this, and the detect route checks the two agree on every call.
    """
    cx, cy = x + fw / 2, y + fh / 2
    side = max(fw, fh) * 2.2
    side = min(side, min(w, h))
    # The face box sets the HEIGHT; the width follows the aspect. Driving it off the
    # height keeps the head the same size in both options, so a 16:9 crop is a square
    # one with more room either side rather than the same picture zoomed out.
    ch = side
    cwid = side * aspect
    if cwid > w:                      # too wide for the frame: give up width, not head
        cwid = w
        ch = min(ch, cwid / aspect)
    # THE CLAMP, and the reason this tool exists. It keeps the crop inside the
    # frame by SLIDING it, so a webcam inset near an edge gets a crop that stays
    # the right size and stops being centred on the face.
    x0 = int(max(0, min(w - cwid, cx - cwid / 2)))
    y0 = int(max(0, min(h - ch, cy - ch / 2)))
    return x0, y0, int(cwid), int(ch), side


def detect_best(video, t0, window, step):
    """The largest, most confident face near t0, with the frame it was found on."""
    det = cv2.FaceDetectorYN.create(MODEL, "", (320, 320), 0.7, 0.3, 5000)
    best = None
    with tempfile.TemporaryDirectory() as td:
        for t in frames(video, t0, window, step):
            p = os.path.join(td, f"f{t}.jpg")
            if not grab(video, t, p):
                continue
            img = cv2.imread(p)
            if img is None:
                continue
            h, w = img.shape[:2]
            det.setInputSize((w, h))
            _, faces = det.detect(img)
            if faces is None:
                continue
            for f in faces:
                x, y, fw, fh, score = f[0], f[1], f[2], f[3], f[-1]
                # Prefer a BIG confident face. Jeremy's webcam inset is small in
                # absolute terms but it is the only face in frame; size still
                # breaks ties between a sharp frame and a motion-blurred one.
                q = float(score) * (fw * fh) ** 0.5
                if best is None or q > best[0]:
                    # keep the DECODED FRAME, not its path -- the temp directory
                    # is gone by the time the crop is taken
                    best = (q, float(x), float(y), float(fw), float(fh),
                            float(score), t, img.copy(), w, h)
    return best


def frame_size(video, t):
    """The frame size, decoded rather than probed -- one code path for both."""
    with tempfile.TemporaryDirectory() as td:
        p = os.path.join(td, "f.jpg")
        if not grab(video, t, p):
            return None, None
        img = cv2.imread(p)
        if img is None:
            return None, None
        return img, img.shape[:2]


def mode_detect(video, t0, window, step, aspect):
    """Report what the detector found. Writes nothing."""
    best = detect_best(video, t0, window, step)
    if best is None:
        _img, size = frame_size(video, t0)
        h, w = size if size else (0, 0)
        print(json.dumps({"face": None, "auto": None, "frame": {"w": int(w), "h": int(h)}}))
        return 1
    _, x, y, fw, fh, score, t, img, w, h = best
    x0, y0, cwid, ch, _side = auto_crop(x, y, fw, fh, w, h, aspect)
    print(json.dumps({
        "face": {"x": x, "y": y, "w": fw, "h": fh, "score": round(float(score), 4), "t": t},
        "auto": {"x": x0, "y": y0, "w": cwid, "h": ch},
        "frame": {"w": int(w), "h": int(h)},
    }))
    return 0


def mode_crop(spec, video, t0, out, size, aspect):
    """Cut exactly the box a person chose, at exactly the frame they chose it on."""
    try:
        x0, y0, cwid, ch = [int(round(float(v))) for v in spec.split(",")]
    except ValueError:
        print("NONE")
        return 1
    img, dims = frame_size(video, t0)
    if img is None:
        print("NONE")
        return 1
    h, w = dims
    # Clamped, not refused: the caller validated this box against the frame, and
    # a numpy slice that runs off the edge returns a smaller picture in silence
    # rather than an error. Better to cut what was asked for as closely as the
    # frame allows and print what was actually cut.
    cwid = max(1, min(w, cwid))
    ch = max(1, min(h, ch))
    x0 = max(0, min(w - cwid, x0))
    y0 = max(0, min(h - ch, y0))
    crop = img[y0:y0 + ch, x0:x0 + cwid]
    crop = cv2.resize(crop, (size, max(1, int(round(size / aspect)))), interpolation=cv2.INTER_AREA)
    cv2.imwrite(out, crop, [cv2.IMWRITE_JPEG_QUALITY, 95])
    # -1.000 where a detector confidence would go: this crop was CHOSEN.
    print(f"{x0} {y0} {ch} -1.000 {t0} {cwid} {ch}")
    return 0


def main():
    # Both flags are ADDITIVE. The positional call is documented in the
    # thumbnails spec sheet and used by make-thumb.mjs, so it is parsed exactly
    # as it always was once the flags are pulled out.
    argv, detect_only, crop_spec = [], False, None
    i = 1
    while i < len(sys.argv):
        a = sys.argv[i]
        if a == "--detect":
            detect_only = True
            i += 1
        elif a == "--crop":
            crop_spec = sys.argv[i + 1]
            i += 2
        else:
            argv.append(a)
            i += 1

    if detect_only:
        video, t0 = argv[0], float(argv[1])
        window = float(argv[2]) if len(argv) > 2 else 0.8
        step = float(argv[3]) if len(argv) > 3 else 0.4
        aspect = float(argv[4]) if len(argv) > 4 else 1.0
        return mode_detect(video, t0, window, step, aspect)

    video, t0, out, size = argv[0], float(argv[1]), argv[2], int(argv[3])

    if crop_spec is not None:
        aspect = float(argv[4]) if len(argv) > 4 else 1.0
        return mode_crop(crop_spec, video, t0, out, size, aspect)

    window = float(argv[4]) if len(argv) > 4 else 0.8
    step = float(argv[5]) if len(argv) > 5 else 0.4
    aspect = float(argv[6]) if len(argv) > 6 else 1.0

    best = detect_best(video, t0, window, step)
    if best is None:
        print("NONE")
        return 1
    _, x, y, fw, fh, score, t, img, w, h = best
    # crop centred on the face, padded out so it is a portrait not a close-up of a
    # nose; clamped to stay inside the frame -- and the clamp is what the face
    # judger exists to make visible, so the maths is auto_crop()'s alone now.
    x0, y0, cwid, ch, side = auto_crop(x, y, fw, fh, w, h, aspect)
    crop = img[y0:y0 + ch, x0:x0 + cwid]
    crop = cv2.resize(crop, (size, max(1, int(round(size / aspect)))), interpolation=cv2.INTER_AREA)
    cv2.imwrite(out, crop, [cv2.IMWRITE_JPEG_QUALITY, 95])
    # The first five fields are the original contract. cwid/ch are appended so
    # make-thumb.mjs can RECORD the box rather than infer it from `side`, which
    # is only the same number while the aspect is square.
    print(f"{x0} {y0} {int(side)} {score:.3f} {t} {cwid} {ch}")
    return 0

if __name__ == "__main__":
    sys.exit(main())