#!/usr/bin/env python3
"""
Two-video drone footage -> 3DGS pipeline for Apple Silicon Macs.

Usage:
    python drone_3dgs_pipeline.py \
        --video1 path/to/flight1.mp4 \
        --video2 path/to/flight2.mp4 \
        --output_dir ./my_scene \
        --fps 2

What it does:
    1. Extracts frames from both videos at the given fps (default 2 fps).
    2. Pools them into one folder with non-colliding names.
    3. Runs COLMAP feature extraction, matching, and sparse reconstruction.
    4. Hands you a Nerfstudio-ready folder structure to train splatfacto.

Then run:
    ns-train splatfacto --data ./my_scene
"""

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


def run(cmd, check=True):
    """Run a shell command, streaming output."""
    print(f"\n>>> {' '.join(str(c) for c in cmd)}\n")
    result = subprocess.run(cmd, check=check)
    return result


def check_dependencies():
    """Verify ffmpeg and colmap are installed."""
    for tool in ("ffmpeg", "colmap"):
        if shutil.which(tool) is None:
            sys.exit(f"ERROR: {tool} not found. Install with: brew install {tool}")


def extract_frames(video_path: Path, out_dir: Path, fps: float, prefix: str):
    """Extract frames from a video using ffmpeg at the given fps."""
    out_dir.mkdir(parents=True, exist_ok=True)
    pattern = str(out_dir / f"{prefix}_%05d.jpg")
    run([
        "ffmpeg", "-i", str(video_path),
        "-r", str(fps),
        "-q:v", "2",          # JPEG quality (2 = high)
        "-y",                  # overwrite
        pattern,
    ])
    count = len(list(out_dir.glob(f"{prefix}_*.jpg")))
    print(f"Extracted {count} frames from {video_path.name} -> {out_dir}")
    return count


def run_colmap(workspace: Path, images_dir: Path, use_gpu: bool = False):
    """Run the COLMAP SfM pipeline: feature extraction, matching, mapping."""
    sparse_dir = workspace / "sparse"
    sparse_dir.mkdir(parents=True, exist_ok=True)
    db_path = workspace / "database.db"

    # Step 1: feature extraction (COLMAP 4.x renamed SiftExtraction → FeatureExtraction)
    run([
        "colmap", "feature_extractor",
        "--database_path", str(db_path),
        "--image_path", str(images_dir),
        "--ImageReader.single_camera", "1",
        "--ImageReader.camera_model", "OPENCV",
        "--FeatureExtraction.use_gpu", "1" if use_gpu else "0",
    ])

    # Step 2: Python matcher (COLMAP 4.x exhaustive_matcher segfaults on Apple Silicon ARM64;
    # match_features.py reads SIFT descriptors from the DB and matches via OpenCV BFMatcher)
    script = Path(__file__).parent / "match_features.py"
    run([sys.executable, str(script), "--db", str(db_path)])

    # Step 3: sparse reconstruction (this is the slow one)
    run([
        "colmap", "mapper",
        "--database_path", str(db_path),
        "--image_path", str(images_dir),
        "--output_path", str(sparse_dir),
    ])

    # COLMAP writes to sparse/0/ by default
    print(f"\nCOLMAP done. Reconstruction in {sparse_dir}/0/")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--video1", type=Path, required=True)
    parser.add_argument("--video2", type=Path, required=True)
    parser.add_argument("--output_dir", type=Path, required=True)
    parser.add_argument("--fps", type=float, default=2.0,
                        help="Frames per second to extract (default: 2). "
                             "Higher = more frames, slower training, better quality.")
    parser.add_argument("--use_gpu", action="store_true",
                        help="Try GPU SIFT in COLMAP (often unreliable on M1; default off).")
    args = parser.parse_args()

    check_dependencies()

    if not args.video1.exists() or not args.video2.exists():
        sys.exit("ERROR: one or both video files not found.")

    workspace = args.output_dir
    images_dir = workspace / "images"
    workspace.mkdir(parents=True, exist_ok=True)

    # Extract frames from both videos into the SAME images folder, with prefixes
    # so filenames don't collide. COLMAP treats them as one set automatically.
    print(f"\n=== Extracting frames from {args.video1.name} ===")
    n1 = extract_frames(args.video1, images_dir, args.fps, prefix="v1")
    print(f"\n=== Extracting frames from {args.video2.name} ===")
    n2 = extract_frames(args.video2, images_dir, args.fps, prefix="v2")
    total = n1 + n2

    print(f"\n=== Total frames: {total} ===")
    if total > 800:
        print("WARNING: lots of frames. Consider lowering --fps. "
              "Exhaustive matching will be slow; switch to sequential_matcher if needed.")

    # Run COLMAP
    print("\n=== Running COLMAP (this is the slow part, get a coffee) ===")
    run_colmap(workspace, images_dir, use_gpu=args.use_gpu)

    # Print next steps
    print(f"""
========================================================================
DONE with SfM. Your scene is at: {workspace}

Next, train the splat. Two options on Mac:

OPTION 1: Nerfstudio (Python, scriptable)
    pip install nerfstudio
    ns-process-data images --data {images_dir} --output-dir {workspace}/ns_data \\
        --skip-colmap --colmap-model-path {workspace}/sparse/0
    ns-train splatfacto --data {workspace}/ns_data

OPTION 2: Brush (Rust binary, faster on Mac)
    Download from https://github.com/ArthurBrussee/brush
    brush --source {workspace}

View result:
    - Online: https://playcanvas.com/supersplat/editor (drag .ply in)
    - Local: install SuperSplat or use Nerfstudio's viewer
========================================================================
""")


if __name__ == "__main__":
    main()