#!/usr/bin/env python3
"""Install the sparse FP4 checkpoint-700 inference overlay into FastVideo.

The checkpoint depends on local FastVideo attention backend changes that are
not part of a vanilla install.  This helper copies the bundled overlay files
into a FastVideo source checkout or site-packages installation.
"""

from __future__ import annotations

import argparse
import importlib.util
import shutil
import sys
from pathlib import Path


def _find_fastvideo_root() -> Path:
    spec = importlib.util.find_spec("fastvideo")
    if spec is None or spec.origin is None:
        raise RuntimeError(
            "Could not import fastvideo. Pass --fastvideo-root explicitly or "
            "activate a FastVideo environment first.")
    return Path(spec.origin).resolve().parents[1]


def _iter_overlay_files(overlay_root: Path):
    for path in sorted(overlay_root.rglob("*")):
        if path.is_file() and "__pycache__" not in path.parts:
            yield path


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--fastvideo-root",
        type=Path,
        default=None,
        help="FastVideo repository/install root. Defaults to import location.",
    )
    parser.add_argument(
        "--backup",
        action="store_true",
        help="Write .sfp4_backup copies before overwriting existing files.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print files that would be copied without modifying anything.",
    )
    args = parser.parse_args()

    bundle_root = Path(__file__).resolve().parent
    overlay_root = bundle_root / "overlay_files"
    if not overlay_root.is_dir():
        raise RuntimeError(f"Missing overlay directory: {overlay_root}")

    target_root = args.fastvideo_root.resolve() if args.fastvideo_root else _find_fastvideo_root()
    if not (target_root / "fastvideo").exists():
        raise RuntimeError(
            f"{target_root} does not look like a FastVideo root: missing fastvideo/")

    copied = 0
    for src in _iter_overlay_files(overlay_root):
        rel = src.relative_to(overlay_root)
        dst = target_root / rel
        print(f"{rel}")
        if args.dry_run:
            continue
        dst.parent.mkdir(parents=True, exist_ok=True)
        if args.backup and dst.exists():
            backup = dst.with_suffix(dst.suffix + ".sfp4_backup")
            if not backup.exists():
                shutil.copy2(dst, backup)
        shutil.copy2(src, dst)
        copied += 1

    if args.dry_run:
        print(f"Dry run complete for target root: {target_root}")
    else:
        print(f"Installed {copied} files into {target_root}")
        print(
            "Use PYTHONPATH='<FastVideo>/fastvideo-kernel/python:"
            "<FastVideo>/fastvideo-kernel:$PYTHONPATH' when running inference.")
    return 0


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