#!/usr/bin/env python3
"""Patch a finite set of replacement frames onto an existing timeline.

This script is intentionally narrow:

- it only replaces frames that exist in the replacement directory
- it does not create a full-timeline cleaned master by itself
- if the base input is the original source, the output will still be source-like
  everywhere outside the listed replacement frames

Use this only after choosing a validated cleaned base, and always run a
whole-timeline audit on the result.
"""

import argparse
import re
from pathlib import Path

import cv2


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=(
            "Apply explicit per-frame replacements onto an existing base video. "
            "This is a patching step, not a full-video cleanup step."
        )
    )
    p.add_argument(
        "--input",
        required=True,
        help=(
            "Base input video to patch. Use a validated full-timeline cleaned base "
            "if you expect the final output to stay cleaned outside the patched spans."
        ),
    )
    p.add_argument(
        "--replacements-dir",
        required=True,
        help=(
            "Directory containing replacement PNG frames named f_<frame>.png. "
            "Only listed frames are replaced."
        ),
    )
    p.add_argument("--output", required=True, help="Output video path")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    repl_dir = Path(args.replacements_dir)
    if not repl_dir.exists():
        raise RuntimeError(f"Replacement dir not found: {repl_dir}")

    mapping: dict[int, Path] = {}
    pat = re.compile(r"f_(\d+)")
    for p in repl_dir.glob("*.png"):
        m = pat.search(p.name)
        if not m:
            continue
        mapping[int(m.group(1))] = p
    if not mapping:
        raise RuntimeError(f"No replacement frames matching 'f_<frameidx>' found in: {repl_dir}")

    cap = cv2.VideoCapture(args.input)
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open input: {args.input}")
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    if out_path.suffix.lower() in {".mkv", ".avi"}:
        fourcc = cv2.VideoWriter_fourcc(*"FFV1")
    else:
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(str(out_path), fourcc, fps, (width, height))
    if not writer.isOpened():
        raise RuntimeError(f"Cannot open output: {out_path}")

    frame_idx = 0
    replaced = 0
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        rp = mapping.get(frame_idx)
        if rp is not None:
            repl = cv2.imread(str(rp))
            if repl is not None and repl.shape[0] == height and repl.shape[1] == width:
                frame = repl
                replaced += 1
        writer.write(frame)
        frame_idx += 1

    writer.release()
    cap.release()
    print(f"Done. frames={frame_idx}, replaced={replaced}, output={out_path}")


if __name__ == "__main__":
    main()