"""Extract evenly spaced JPG frames from a source video into public/frames/.

The Hero canvas scrubs through these frames on scroll, so the sequence must be
evenly spaced across the *whole* video and numbered without gaps.

FFmpeg is not installed system-wide here; imageio-ffmpeg ships a static binary
and pulls in no numpy, which matters because the local numpy is broken.

Frames are scaled straight from the decoded video rather than by re-encoding
the previous JPGs, so an upscale does not bake in the artefacts of the smaller
pass and then compress them a second time.

Usage:
    python scripts/extract_frames.py [--count 120] [--width 3840] [--height 2160]
"""

from __future__ import annotations

import argparse
import shutil
import subprocess
import sys
from pathlib import Path

import imageio_ffmpeg

ROOT = Path(__file__).resolve().parent.parent
VIDEO = ROOT / "Solar_energy.mp4"
OUT_DIR = ROOT / "public" / "frames"


def probe(ffmpeg: str, video: Path) -> tuple[int, float, int, int]:
    """Return (frame_count, duration_secs, width, height)."""
    n_frames, secs = imageio_ffmpeg.count_frames_and_secs(str(video))
    meta = imageio_ffmpeg.read_frames(str(video)).__next__()
    width, height = meta["size"]
    return n_frames, secs, width, height


def build_select_expr(indices: list[int]) -> str:
    # Commas are filter separators in ffmpeg's parser, so they must be escaped
    # inside the expression itself — this is ffmpeg syntax, not shell quoting.
    return "+".join(r"eq(n\,{})".format(i) for i in indices)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--count", type=int, default=120)
    ap.add_argument("--width", type=int, default=3840)
    ap.add_argument("--height", type=int, default=2160)
    ap.add_argument("--quality", type=int, default=1, help="ffmpeg -q:v, 1=best")
    ap.add_argument(
        "--no-sharpen",
        action="store_true",
        help="skip the unsharp pass that offsets upscaling softness",
    )
    ap.add_argument(
        "--delogo",
        default="1698:848:88:100",
        help=(
            "x:y:w:h of the burnt-in watermark in SOURCE pixels, patched by "
            "interpolating from the box border. Pass '' to keep it. These "
            "coordinates are tied to the source resolution, so re-measure "
            "them if --video changes."
        ),
    )
    ap.add_argument("--video", type=Path, default=VIDEO)
    args = ap.parse_args()

    if not args.video.exists():
        print("Video not found: {}".format(args.video), file=sys.stderr)
        return 1

    ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
    total, secs, w, h = probe(ffmpeg, args.video)
    print("Source : {}".format(args.video.name))
    print("Frames : {}  ({:.2f}s, {}x{})".format(total, secs, w, h))

    if total < args.count:
        print(
            "Video has only {} frames; extracting all of them instead of {}.".format(
                total, args.count
            )
        )
    n_out = min(args.count, total)

    # Evenly spaced across the full range, first and last frame inclusive.
    if n_out == 1:
        indices = [0]
    else:
        indices = [round(i * (total - 1) / (n_out - 1)) for i in range(n_out)]
    indices = sorted(set(indices))

    # Start from a clean directory so a shorter run cannot leave stale frames
    # behind that the canvas would happily keep scrubbing through.
    if OUT_DIR.exists():
        shutil.rmtree(OUT_DIR)
    OUT_DIR.mkdir(parents=True)

    vf = "select='{}'".format(build_select_expr(indices))
    if args.delogo:
        # Patch the watermark at source resolution, before the upscale — the
        # filter has real neighbouring pixels to interpolate from here, where
        # after scaling it would only ever see invented ones.
        dx, dy, dw, dh = args.delogo.split(":")
        vf += ",delogo=x={}:y={}:w={}:h={}".format(dx, dy, dw, dh)
        print("Delogo : x={} y={} w={} h={} (source px)".format(dx, dy, dw, dh))
    vf += ",scale={}:{}:flags=lanczos".format(args.width, args.height)
    if not args.no_sharpen and args.width > w:
        # Any upscale softens edges. A light unsharp brings back the panel
        # grid and cable detail without the halos a stronger amount produces.
        vf += ",unsharp=5:5:0.8:5:5:0.0"
    print("Target : {}x{}  (q:v {})".format(args.width, args.height, args.quality))

    cmd = [
        ffmpeg, "-hide_banner", "-loglevel", "error", "-y",
        "-i", str(args.video),
        "-vf", vf,
        "-vsync", "0",
        "-q:v", str(args.quality),
        "-start_number", "1",
        str(OUT_DIR / "frame_%04d.jpg"),
    ]
    print("Extracting {} frames -> {}".format(len(indices), OUT_DIR))
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(result.stderr.strip(), file=sys.stderr)
        return result.returncode

    written = sorted(OUT_DIR.glob("frame_*.jpg"))
    total_mb = sum(f.stat().st_size for f in written) / (1024 * 1024)
    print("Wrote  : {} files, {:.1f} MB total".format(len(written), total_mb))
    if written:
        print("Range  : {} .. {}".format(written[0].name, written[-1].name))

    # A gap in the numbering means the canvas would request a 404 mid-scroll.
    expected = ["frame_{:04d}.jpg".format(i + 1) for i in range(len(written))]
    missing = [n for n in expected if not (OUT_DIR / n).exists()]
    if missing:
        print("Gaps in sequence: {}".format(missing[:5]), file=sys.stderr)
        return 1

    print("FRAME_COUNT = {}".format(len(written)))
    return 0


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