# /// script
# requires-python = ">=3.11"
# dependencies = ["imageio-ffmpeg"]
# ///
"""Extract one thumbnail per Prelinger chunk (frame at the chunk midpoint).

Work list comes from /data/chunkmeta/*.json — the recorded stream-copy
boundaries — so the midpoint is (end - start) / 2 of the ACTUAL segment, no
ffprobe round-trip. Chunk mp4s have reset timestamps, so that offset is also
the seek position inside the file.

In:  /pf/chunks/{id}/c%04d.mp4
     /pf/chunkmeta/{id}.json
Out: /pf/thumbs/{id}/c%04d.jpg
     /pf/thumbmeta/{id}.json   (written LAST = per-film completion marker)
"""

import argparse
import json
import shutil
import subprocess
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import imageio_ffmpeg

FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()  # default uv image ships no ffmpeg
ROOT = Path("/pf")  # /data is reserved for Jobs artifacts on local-script runs
CHUNKS, CMETA = ROOT / "chunks", ROOT / "chunkmeta"
THUMBS, TMETA = ROOT / "thumbs", ROOT / "thumbmeta"
WIDTH = 480
QSCALE = "4"  # mjpeg -q:v 4 ~= quality 80


def exists(p: Path) -> bool:
    # FUSE stat races with concurrent writers; a raised OSError is not "absent"
    # in a way we can act on, but treating it as absent only costs a redo.
    try:
        return p.exists()
    except OSError:
        return False


def grab(src: Path, dst_tmp: Path, at: float) -> bool:
    """One frame at `at` seconds, scaled to WIDTH. -ss before -i = fast seek."""
    try:
        subprocess.run(
            [FFMPEG, "-hide_banner", "-loglevel", "error",
             "-ss", f"{at:.3f}", "-i", str(src),
             "-frames:v", "1", "-vf", f"scale={WIDTH}:-2",
             "-q:v", QSCALE, "-y", str(dst_tmp)],
            check=True, stdin=subprocess.DEVNULL, timeout=180)
    except (subprocess.SubprocessError, OSError):
        return False
    return dst_tmp.exists() and dst_tmp.stat().st_size > 0


def thumb_one(metafile: Path) -> str:
    try:
        meta = json.loads(metafile.read_text())
    except (OSError, json.JSONDecodeError):
        print(f"[err] unreadable chunkmeta {metafile.name}", flush=True)
        return "err"
    ident = meta["identifier"]
    marker = TMETA / f"{ident}.json"
    if exists(marker):
        return "skip"

    thumbs, failed = [], []
    with tempfile.TemporaryDirectory() as td:
        tdp = Path(td)
        dest = THUMBS / ident
        try:
            dest.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            print(f"[err] {ident}: mkdir {type(e).__name__}", flush=True)
            return "err"
        for c in meta["chunks"]:
            name = c["chunk"]
            stem = Path(name).stem
            mid = max(0.0, (float(c["end"]) - float(c["start"])) / 2.0)
            src = CHUNKS / ident / name
            local = tdp / f"{stem}.jpg"
            if not grab(src, local, mid):
                # retry once at the first frame — very short or truncated tails
                if not grab(src, local, 0.0):
                    failed.append(name)
                    continue
                mid = 0.0
            try:
                # FUSE has no rename: write local, then copy onto the mount
                shutil.copyfile(local, dest / f"{stem}.jpg")
            except OSError:
                failed.append(name)
                continue
            thumbs.append({"chunk": name, "thumb": f"{stem}.jpg",
                           "t": round(mid, 3)})

    try:
        marker.write_text(json.dumps({"identifier": ident, "thumbs": thumbs,
                                      "failed": failed}))
    except OSError as e:
        print(f"[err] {ident}: marker {type(e).__name__}", flush=True)
        return "err"
    if failed:
        print(f"[partial] {ident}: {len(failed)} chunk(s) failed", flush=True)
        return "partial"
    return "ok"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=8)
    args = ap.parse_args()

    THUMBS.mkdir(parents=True, exist_ok=True)
    TMETA.mkdir(parents=True, exist_ok=True)
    films = sorted(CMETA.glob("*.json"))
    print(f"[plan] {len(films)} films with chunkmeta, workers={args.workers}",
          flush=True)

    counts = {"ok": 0, "skip": 0, "partial": 0, "err": 0}
    t0 = time.time()
    with ThreadPoolExecutor(args.workers) as ex:
        for n, res in enumerate(ex.map(thumb_one, films), 1):
            counts[res] += 1
            if n % 50 == 0 or n == len(films):
                print(f"[progress] {n}/{len(films)} {counts} "
                      f"{time.time() - t0:.0f}s", flush=True)
    try:
        total = sum(1 for _ in THUMBS.rglob("*.jpg"))
    except OSError:
        total = -1
    print(f"[THUMBS DONE] {counts} total_thumbs={total} "
          f"{time.time() - t0:.0f}s", flush=True)


if __name__ == "__main__":
    main()