"""Harvest still frames from an archive VOD into the ND1 dataset -- design
doc 7_FarZone_Detection_BallID.md section 3.5, step 3 of the BRIEF of
28/08/2026.

WHY. Every fix in that design doc is a patch that buys frames back without
retraining: more pixels through the detector, a low-confidence recovery lane
for tracks that are already alive. The real repair is a detector fine-tuned
on THIS footage -- far balls of 20 px, balls hiding behind balls in a rack
seen down its own axis -- and that needs a labelled set. This script is the
first half of it: pull the frames out, record where every one came from, and
stop. Labelling and training are later sessions.

WHAT IT GUARANTEES.

- The archive is READ ONLY. ``D:\\Khoa luan`` is opened for reading and
  nothing else; the whole tree is fingerprinted (every file's size and
  mtime) before and after and the two are compared, so "no byte was written
  to the archive" is a measurement in the manifest, not a promise.
- Every frame is traceable. The manifest records the source path, the
  SHA-256 of the source file, the nominal fps and duration, and for each
  frame its file name, its index and its container PTS. A frame whose
  provenance is lost is a frame that cannot be used as training data.
- The output lands OUTSIDE both repositories, in
  ``D:\\poolcoach\\datasets\\nd1_frames\\<vod stem>\\``: broadcast footage is
  not committed, and neither are its frames.

SAMPLING. One frame every ``--every`` seconds (2 s by default), taken on the
container's own PTS, which is what makes the count come out at
duration / every within a frame. ``--shots <shots.json>`` optionally ADDS
frames around the moments a shot list already knows about -- the strike and
the settle of every shot, and the seconds before a rack is broken, which is
failure mode 2 of section 1 and the hardest thing in the set to label. The
extra frames are snapped onto the same PTS grid and de-duplicated against
the uniform ones, and the manifest counts the two kinds separately.

RUNNING IT (the launcher owns the paths):

    ..\\poolcoach-cv-env\\Scripts\\python.exe scripts\\harvest_frames.py ^
        --video "D:\\Khoa luan\\...\\rack_5245_5805.mp4" ^
        --out   "D:\\poolcoach\\datasets\\nd1_frames"
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
for _p in (ROOT / "src", ROOT):
    if str(_p) not in sys.path:
        sys.path.insert(0, str(_p))

EVERY_S = 2.0             # the default sampling period. A rack lasts minutes
                          # and the table barely changes between two
                          # neighbouring frames; 1 frame / 2 s is what
                          # section 3.5 asks for
JPEG_QUALITY = 95         # training data, not a thumbnail: the artefacts of
                          # a cheap JPEG land on exactly the 20 px balls the
                          # set exists to teach
SHOT_OFFSETS_S = (-1.5, -0.5, 0.0, 0.3, 0.8)
                          # around each shot's onset when a shot list is
                          # given: the rack still standing, the strike, and
                          # the first moments of the spread -- the states a
                          # uniform grid hits only by luck
HASH_CHUNK = 1 << 20


def sha256_of(path: Path) -> str:
    """The source file's SHA-256 -- the only way a frame harvested today can
    be tied to the exact file it came from tomorrow."""
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        while True:
            chunk = fh.read(HASH_CHUNK)
            if not chunk:
                break
            h.update(chunk)
    return h.hexdigest()


def fingerprint(root: Path) -> dict:
    """Every file under ``root`` as name -> (size, mtime_ns), so "the
    archive was not written to" can be CHECKED after the run instead of
    assumed."""
    out = {}
    for p in sorted(root.rglob("*")):
        try:
            st = p.stat()
        except OSError:
            continue
        if p.is_file():
            out[str(p.relative_to(root))] = [st.st_size, st.st_mtime_ns]
    return out


def wanted_times(duration_s: float, every: float,
                 shots_path: Path | None) -> tuple[list[float], set[float]]:
    """The timestamps to harvest: the uniform grid, plus the shot-relative
    ones when a shot list is given.

    Returns (sorted times, the subset that came from the shot list). The
    uniform grid is ``0, every, 2*every, ...`` strictly below the duration,
    which is what makes the count land on ``duration / every`` within one
    frame -- the acceptance number of step 3.
    """
    grid: list[float] = []
    t = 0.0
    while t < duration_s:
        grid.append(round(t, 3))
        t += every
    extra: set[float] = set()
    if shots_path is not None:
        manifest = json.loads(shots_path.read_text(encoding="utf-8"))
        for shot in manifest.get("shots", []):
            base = shot.get("t_onset_s")
            if base is None:
                continue
            for off in SHOT_OFFSETS_S:
                t = round(float(base) + off, 3)
                if 0.0 <= t < duration_s:
                    extra.add(t)
        extra -= set(grid)
    return sorted(set(grid) | extra), extra


def harvest(video: Path, out_dir: Path, every: float,
            shots_path: Path | None, limit: int | None) -> dict:
    """Decode the VOD ONCE from start to end, keeping the frame nearest each
    wanted timestamp.

    Sequential decoding rather than seeking: a seek on a long-GOP broadcast
    file lands on the previous keyframe, so the PTS written into the
    manifest would not be the PTS of the image saved next to it. Decoding
    320 s costs seconds and buys an honest manifest.
    """
    import cv2

    cap = cv2.VideoCapture(str(video))
    if not cap.isOpened():
        raise SystemExit(f"cannot open {video}")
    try:
        fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
        n_est = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0
        duration = n_est / fps if fps > 0 and n_est > 0 else 0.0
        if duration <= 0:
            raise SystemExit(f"{video} reports no duration (fps={fps}, "
                             f"frames={n_est}) -- it cannot be sampled")
        times, from_shots = wanted_times(duration, every, shots_path)
        if limit:
            times = times[:limit]
        out_dir.mkdir(parents=True, exist_ok=True)

        frames: list[dict] = []
        idx = -1
        want = 0
        t0 = time.perf_counter()
        while want < len(times):
            if not cap.grab():
                break
            idx += 1
            t_s = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0
            if t_s + 1e-9 < times[want]:
                continue
            ok, frame = cap.retrieve()
            if not ok:
                break
            name = f"{video.stem}_{idx:06d}.jpg"
            cv2.imwrite(str(out_dir / name), frame,
                        [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
            frames.append({"file": name, "frame_idx": idx,
                           "t_s": round(t_s, 3),
                           "t_wanted_s": times[want],
                           "source": "shots" if times[want] in from_shots
                           else "grid"})
            want += 1
            if len(frames) % 25 == 0:
                print(f"[harvest] {len(frames)}/{len(times)} frames "
                      f"(t={t_s:.1f}s)", flush=True)
        elapsed = time.perf_counter() - t0
    finally:
        cap.release()

    return {"fps_nominal": round(fps, 3), "duration_s": round(duration, 3),
            "n_frames_video": int(n_est), "every_s": every,
            "n_wanted": len(times), "n_written": len(frames),
            "n_from_shots": sum(1 for f in frames if f["source"] == "shots"),
            "elapsed_s": round(elapsed, 1), "frames": frames}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(
        description="Harvest ND1 training frames from an archive VOD "
                    "(design doc 7 section 3.5). The archive is read only.")
    ap.add_argument("--video", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path,
                    help="the ND1 root; the frames go in <out>/<vod stem>/")
    ap.add_argument("--every", type=float, default=EVERY_S,
                    help=f"seconds between frames (default {EVERY_S:g})")
    ap.add_argument("--shots", type=Path, default=None,
                    help="an existing shots.json -- ADDS frames around every "
                         "shot onset (the rack before the break included)")
    ap.add_argument("--limit", type=int, default=None,
                    help="stop after N frames (a smoke run)")
    ap.add_argument("--archive-root", type=Path, default=None,
                    help="the read-only tree to fingerprint before and after "
                         "(default: the video's own drive root is NOT "
                         "guessed -- pass it explicitly to get the check)")
    args = ap.parse_args(argv)

    if not args.video.exists():
        raise SystemExit(f"{args.video} does not exist")
    out_dir = args.out / args.video.stem
    print(f"[harvest] video  : {args.video} "
          f"({args.video.stat().st_size / 1e6:.1f} MB)")
    print(f"[harvest] out    : {out_dir}", flush=True)

    fp_before = None
    if args.archive_root is not None:
        print(f"[harvest] fingerprinting {args.archive_root} ...", flush=True)
        fp_before = fingerprint(args.archive_root)
        print(f"[harvest] {len(fp_before)} files", flush=True)

    print("[harvest] hashing the source ...", flush=True)
    digest = sha256_of(args.video)
    print(f"[harvest] sha256 {digest}", flush=True)

    result = harvest(args.video, out_dir, args.every, args.shots, args.limit)

    archive = None
    if fp_before is not None:
        fp_after = fingerprint(args.archive_root)
        changed = sorted(k for k in set(fp_before) | set(fp_after)
                         if fp_before.get(k) != fp_after.get(k))
        archive = {"root": str(args.archive_root), "n_files": len(fp_before),
                   "n_changed": len(changed), "changed": changed[:20]}

    manifest = {
        "dataset": "ND1",
        "design_doc": "7_FarZone_Detection_BallID.md section 3.5",
        "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "source": {"path": str(args.video), "name": args.video.name,
                   "size_bytes": args.video.stat().st_size,
                   "sha256": digest},
        "sampling": {"every_s": args.every,
                     "shots": str(args.shots) if args.shots else None,
                     "shot_offsets_s": list(SHOT_OFFSETS_S)
                     if args.shots else None,
                     "limit": args.limit},
        "jpeg_quality": JPEG_QUALITY,
        "fps_nominal": result["fps_nominal"],
        "duration_s": result["duration_s"],
        "n_frames_video": result["n_frames_video"],
        "n_wanted": result["n_wanted"],
        "n_written": result["n_written"],
        "n_from_shots": result["n_from_shots"],
        "elapsed_s": result["elapsed_s"],
        "archive_check": archive,
        "frames": result["frames"],
    }
    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=1),
                                           encoding="utf-8")

    expected = result["duration_s"] / args.every
    print(f"[harvest] {result['n_written']} frames written in "
          f"{result['elapsed_s']:.1f}s "
          f"({result['n_from_shots']} of them around a shot onset)")
    print(f"[harvest] duration/every = {expected:.1f} -- written "
          f"{result['n_written']}")
    print(f"[harvest] manifest {out_dir / 'manifest.json'}")
    if archive is not None:
        if archive["n_changed"]:
            print(f"[harvest] RED: {archive['n_changed']} file(s) under "
                  f"{archive['root']} changed -- {archive['changed'][:5]}")
            return 7
        print(f"[harvest] archive untouched ({archive['n_files']} files "
              f"identical)")
    else:
        print("[harvest] NOTE: no --archive-root given, so nothing verified "
              "that the source tree was left alone")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())