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

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

if not a.input:
    print("ERROR: no input frames folder")
    sys.exit(1)
src = Path(a.input[0])
if not src.is_dir():
    print(f"ERROR: input must be a frames folder, got {src}")
    sys.exit(1)

outdir = Path(a.outdir)
dataset = outdir / "colmap_dataset"
images = dataset / "images"
images.mkdir(parents=True, exist_ok=True)
frames = sorted([f for f in src.iterdir() if f.suffix.lower() in (".jpg", ".jpeg", ".png")])
if len(frames) < 3:
    print(f"ERROR: need >=3 frames, found {len(frames)}")
    sys.exit(1)
for f in frames:
    shutil.copy2(f, images / f.name)
print(f"copied {len(frames)} frames", flush=True)

db_path = dataset / "database.db"
sparse = dataset / "sparse"
sparse.mkdir(exist_ok=True)


def run(*cmd):
    print("+", " ".join(str(c) for c in cmd), flush=True)
    r = subprocess.run([str(c) for c in cmd], stdout=sys.stdout, stderr=subprocess.STDOUT)
    if r.returncode != 0:
        print(f"ERROR: command failed ({r.returncode})")
        sys.exit(r.returncode)


run("colmap", "feature_extractor", "--database_path", db_path, "--image_path", images,
    "--ImageReader.single_camera", "1" if p.get("single_camera", True) else "0",
    "--ImageReader.camera_model", p.get("camera_model", "OPENCV"),
    "--FeatureExtraction.use_gpu", "0")

matcher = "sequential_matcher" if p.get("matcher", "sequential") == "sequential" else "exhaustive_matcher"
run("colmap", matcher, "--database_path", db_path, "--FeatureMatching.use_gpu", "0")

if p.get("mapper", "global") == "global":
    run("colmap", "global_mapper", "--database_path", db_path,
        "--image_path", images, "--output_path", sparse)
else:
    run("colmap", "mapper", "--database_path", db_path,
        "--image_path", images, "--output_path", sparse)

# report registered images
recon = sparse / "0"
n_reg = 0
if (recon / "images.bin").exists():
    n_reg = (recon / "images.bin").stat().st_size  # crude presence signal
print(f"reconstruction at {recon}, images.bin present: {(recon / 'images.bin').exists()}", flush=True)
if not (recon / "images.bin").exists() and not (recon / "images.txt").exists():
    print("ERROR: COLMAP produced no reconstruction (too few features/matches?)")
    sys.exit(1)

(outdir / "result.json").write_text(json.dumps({"outputs": [
    {"path": str(dataset), "name": f"{src.name}_colmap",
     "meta": {"kind": "colmap_dataset", "frames": len(frames)}}]}))
print("done", flush=True)