"""COLMAP structure-from-motion preprocessing (runs inside the Space, on CPU).

Takes a set of raw multi-view images of a scene, estimates camera intrinsics /
poses / a sparse point cloud with COLMAP (via `pycolmap`), undistorts the images
to a PINHOLE camera model and writes a standard COLMAP scene layout

    <out>/images/*.jpg
    <out>/sparse/{cameras,images,points3D}.bin

which is exactly what CodecSplat's reference COLMAP loader expects.
"""

from __future__ import annotations

import shutil
from pathlib import Path

import pycolmap
from PIL import Image, ImageOps

# ZeroGPU hosts expose ~192 CPUs. With COLMAP's default num_threads=-1 the incremental
# mapper then fails to register anything beyond the initial pair (reproduced locally by
# forcing 192 threads), so cap the thread count explicitly.
NUM_THREADS = 16

IMAGE_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff")


def prepare_images(image_paths, dest: Path, max_image_size: int = 1280) -> list[str]:
    """Copy the uploaded images into `dest`, EXIF-rotated, RGB, size-capped."""
    dest.mkdir(parents=True, exist_ok=True)
    names: list[str] = []
    for i, src in enumerate(str(p) for p in image_paths):  # caller decides the order
        if not src.lower().endswith(IMAGE_EXT):
            continue
        try:
            img = Image.open(src)
            img = ImageOps.exif_transpose(img).convert("RGB")
        except Exception as exc:
            print(f"[colmap] skipping unreadable image {src}: {exc}", flush=True)
            continue
        scale = max_image_size / max(img.size)
        if scale < 1.0:
            img = img.resize(
                (max(1, int(round(img.width * scale))), max(1, int(round(img.height * scale)))),
                Image.LANCZOS,
            )
        name = f"{i:04d}.jpg"
        img.save(dest / name, quality=95)
        names.append(name)
    return names


def _largest_reconstruction(maps) -> pycolmap.Reconstruction | None:
    recs = list(maps.values()) if isinstance(maps, dict) else list(maps)
    if not recs:
        return None
    return max(recs, key=lambda r: r.num_reg_images())


def run_sfm(
    image_paths,
    workdir: Path,
    max_image_size: int = 1280,
    matcher: str = "exhaustive",
    max_num_features: int = 8192,
    log=print,
) -> dict:
    """Run COLMAP SfM on `image_paths`; return info about the written scene.

    Returns a dict with `scene_dir`, `sparse_sub`, `images_sub`, `num_input`,
    `num_registered`, `num_points3D`, `camera_model`.
    """
    workdir = Path(workdir)
    image_dir = workdir / "sfm_images"
    names = prepare_images(image_paths, image_dir, max_image_size)
    if len(names) < 3:
        raise ValueError(
            f"COLMAP needs at least 3 usable images, got {len(names)}. "
            "Upload several overlapping photos of the same scene."
        )
    log(f"[colmap] {len(names)} images prepared at <= {max_image_size}px")
    try:
        import os as _os
        cpu = next((l.split(":", 1)[1].strip() for l in open("/proc/cpuinfo") if l.startswith("model name")), "?")
        log(f"[colmap] cpu={cpu} ncpu={_os.cpu_count()} threads={NUM_THREADS}")
    except Exception:
        pass

    database_path = workdir / "colmap.db"
    if database_path.exists():
        database_path.unlink()

    extraction_options = pycolmap.FeatureExtractionOptions()
    extraction_options.sift.max_num_features = int(max_num_features)
    extraction_options.sift.estimate_affine_shape = False
    extraction_options.sift.domain_size_pooling = False
    extraction_options.use_gpu = False
    extraction_options.max_image_size = int(max_image_size)
    extraction_options.num_threads = NUM_THREADS

    reader_options = pycolmap.ImageReaderOptions()
    reader_options.camera_model = "SIMPLE_RADIAL"

    pycolmap.extract_features(
        database_path=database_path,
        image_path=image_dir,
        camera_mode=pycolmap.CameraMode.SINGLE,
        reader_options=reader_options,
        extraction_options=extraction_options,
        device=pycolmap.Device.cpu,
    )
    log("[colmap] features extracted")

    matching_options = pycolmap.FeatureMatchingOptions()
    matching_options.num_threads = NUM_THREADS
    if matcher == "sequential":
        pairing = pycolmap.SequentialPairingOptions()
        pairing.overlap = 10
        pairing.quadratic_overlap = True
        pycolmap.match_sequential(
            database_path=database_path,
            matching_options=matching_options,
            pairing_options=pairing,
            device=pycolmap.Device.cpu,
        )
    else:
        pycolmap.match_exhaustive(database_path=database_path, matching_options=matching_options,
                                  device=pycolmap.Device.cpu)
    log(f"[colmap] {matcher} matching done")

    sparse_dir = workdir / "sparse_raw"
    sparse_dir.mkdir(parents=True, exist_ok=True)
    options = pycolmap.IncrementalPipelineOptions()
    options.min_num_matches = 15
    options.num_threads = NUM_THREADS
    options.multiple_models = False
    options.extract_colors = True
    maps = pycolmap.incremental_mapping(
        database_path=database_path,
        image_path=image_dir,
        output_path=sparse_dir,
        options=options,
    )
    rec = _largest_reconstruction(maps)
    if rec is None or rec.num_reg_images() < 3:
        raise ValueError(
            "COLMAP could not reconstruct this image set (only "
            f"{0 if rec is None else rec.num_reg_images()} of {len(names)} images registered). "
            "Use photos with plenty of visual overlap between neighbours, avoid pure "
            "rotation / blurry frames, and keep the scene static."
        )
    # Deterministic world scale (same convention COLMAP itself uses), so that the
    # fixed near/far planes of the CodecSplat checkpoint stay meaningful.
    rec.normalize(fixed_scale=False, extent=10.0)
    log(f"[colmap] mapped {rec.num_reg_images()}/{len(names)} images, {rec.num_points3D()} points")

    model_dir = workdir / "sparse_model"
    if model_dir.exists():
        shutil.rmtree(model_dir)
    model_dir.mkdir(parents=True)
    rec.write(str(model_dir))

    scene_dir = workdir / "colmap_scene"
    if scene_dir.exists():
        shutil.rmtree(scene_dir)
    pycolmap.undistort_images(
        output_path=scene_dir,
        input_path=model_dir,
        image_path=image_dir,
    )
    log("[colmap] undistorted to PINHOLE")

    # Keep only what CodecSplat reads (the undistorter also writes stereo/ scaffolding).
    for extra in ("stereo",):
        p = scene_dir / extra
        if p.is_dir():
            shutil.rmtree(p, ignore_errors=True)

    cam_model = "PINHOLE"
    try:
        cam_model = list(rec.cameras.values())[0].model.name
    except Exception:
        pass

    # Read back the undistorted PINHOLE cameras (normalized K, OpenCV c2w), sorted by name.
    und = pycolmap.Reconstruction(str(scene_dir / "sparse"))
    frames = []
    for _, img in und.images.items():
        cam = und.cameras[img.camera_id]
        p = list(cam.params)
        if cam.model.name == "PINHOLE":
            fx, fy, cx, cy = p[:4]
        elif cam.model.name == "SIMPLE_PINHOLE":
            fx = fy = p[0]
            cx, cy = p[1:3]
        else:
            continue
        cfw = img.cam_from_world() if callable(img.cam_from_world) else img.cam_from_world
        w2c = cfw.matrix()  # 3x4
        path = scene_dir / "images" / img.name
        if path.exists():
            frames.append({
                "name": img.name, "path": str(path),
                "K": [[fx / cam.width, 0, cx / cam.width], [0, fy / cam.height, cy / cam.height], [0, 0, 1]],
                "w2c": [list(map(float, row)) for row in w2c],
            })
    frames.sort(key=lambda f: f["name"])

    return {
        "frames": frames,
        "scene_dir": str(scene_dir),
        "sparse_sub": "sparse",
        "images_sub": "images",
        "num_input": len(names),
        "num_registered": int(rec.num_reg_images()),
        "num_points3D": int(rec.num_points3D()),
        "camera_model": cam_model,
    }


def run_sfm_subprocess(image_paths, workdir, max_image_size=1280, matcher="exhaustive",
                       timeout=900, log=print) -> dict:
    """Run `run_sfm` in a fresh Python process that only imports pycolmap.

    Inside the Gradio process torch / CUDA / OpenMP / BLAS libraries are already loaded,
    which makes pycolmap's bundled OpenBLAS / OpenMP misbehave (crashes in matching and
    failed image registration). A clean child process avoids all of that.
    """
    import json
    import os
    import subprocess
    import sys

    workdir = Path(workdir)
    workdir.mkdir(parents=True, exist_ok=True)
    spec = workdir / "sfm_job.json"
    result = workdir / "sfm_result.json"
    spec.write_text(json.dumps({
        "image_paths": [str(p) for p in image_paths], "workdir": str(workdir),
        "max_image_size": int(max_image_size), "matcher": matcher, "result": str(result),
    }))
    env = {k: v for k, v in os.environ.items() if k not in ("LD_PRELOAD",)}
    env["LD_LIBRARY_PATH"] = os.environ.get("_ORIG_LD_LIBRARY_PATH", "")
    env["OPENBLAS_NUM_THREADS"] = "1"
    env.pop("CUDA_VISIBLE_DEVICES", None)
    proc = subprocess.run([sys.executable, os.path.abspath(__file__), str(spec)], env=env,
                          capture_output=True, text=True, timeout=timeout)
    for line in (proc.stdout or "").splitlines():
        if line.startswith("[colmap]"):
            log(line)
    if not result.exists():
        tail = "\n".join(((proc.stderr or "") + (proc.stdout or "")).splitlines()[-15:])
        raise RuntimeError(f"COLMAP subprocess failed (exit {proc.returncode}):\n{tail}")
    data = json.loads(result.read_text())
    if "error" in data:
        raise ValueError(data["error"])
    return data


if __name__ == "__main__":
    import json
    import sys

    job = json.loads(Path(sys.argv[1]).read_text())
    try:
        out = run_sfm(job["image_paths"], Path(job["workdir"]), job["max_image_size"],
                      job["matcher"], log=lambda m: print(m, flush=True))
    except ValueError as exc:
        out = {"error": str(exc)}
    Path(job["result"]).write_text(json.dumps(out))