"""Camera downloads in Nerfstudio/OpenGL and COLMAP/OpenCV conventions."""

import json
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile

import numpy as np
import trimesh
from scipy.spatial.transform import Rotation


def export_cameras(root: Path, c2w: np.ndarray, intrinsics: np.ndarray,
                   image_names: list[str], width: int, height: int) -> tuple[str, str]:
    """Export poses for the processed images without changing the world frame."""
    if len(c2w) != len(intrinsics) or len(c2w) != len(image_names):
        raise ValueError("Camera and image counts differ")
    if not np.isfinite(c2w).all() or not np.isfinite(intrinsics).all():
        raise ValueError("Cameras must be finite")
    frames, cameras, images = [], [], []
    flip = np.diag([1.0, -1.0, -1.0, 1.0])
    for index, (pose, k, name) in enumerate(zip(c2w, intrinsics, image_names), 1):
        if Path(name).name != name or any(ch.isspace() for ch in name):
            raise ValueError("Export filenames must be simple basenames")
        fx, fy, cx, cy = k[0, 0], k[1, 1], k[0, 2], k[1, 2]
        frames.append({
            "file_path": f"images/{name}",
            "transform_matrix": (pose @ flip).tolist(),
            "fl_x": float(fx), "fl_y": float(fy),
            "cx": float(cx), "cy": float(cy), "w": width, "h": height,
        })
        cameras.append(f"{index} PINHOLE {width} {height} {fx:.12g} {fy:.12g} {cx:.12g} {cy:.12g}\n")
        w2c = np.linalg.inv(pose)
        xyzw = Rotation.from_matrix(w2c[:3, :3]).as_quat()
        qwxyz = xyzw[[3, 0, 1, 2]]
        values = " ".join(f"{value:.12g}" for value in np.r_[qwxyz, w2c[:3, 3]])
        # COLMAP requires a second (empty here) observations line per image.
        images.append(f"{index} {values} {index} {name}\n\n")
    transforms = root / "transforms.json"
    transforms.write_text(json.dumps({"camera_model": "OPENCV", "k1": 0.0,
        "k2": 0.0, "p1": 0.0, "p2": 0.0, "frames": frames}, indent=2) + "\n")
    sparse = root / "sparse" / "0"
    sparse.mkdir(parents=True, exist_ok=True)
    (sparse / "cameras.txt").write_text("# CAMERA_ID MODEL WIDTH HEIGHT PARAMS[]\n" + "".join(cameras))
    (sparse / "images.txt").write_text("# IMAGE_ID QW QX QY QZ TX TY TZ CAMERA_ID NAME\n# Next line: POINTS2D[] (empty)\n" + "".join(images))
    (sparse / "points3D.txt").write_text("# Camera-only model; no reconstructed 3D points.\n")
    outputs = []
    for filename, metadata in [("nerfstudio.zip", [transforms]),
                               ("colmap.zip", list(sparse.glob("*.txt")))]:
        archive = root / filename
        with ZipFile(archive, "w", compression=ZIP_DEFLATED) as handle:
            for path in metadata + [root / "images" / name for name in image_names]:
                handle.write(path, path.relative_to(root).as_posix())
        outputs.append(str(archive))
    return tuple(outputs)


def camera_scene(root: Path, c2w: np.ndarray, intrinsics: np.ndarray,
                 width: int, height: int, trajectory: np.ndarray | None = None) -> str:
    """Create a GLB of camera frusta and optional render path (no inferred mesh)."""
    scene = trimesh.Scene()
    centers = c2w[:, :3, 3]
    span = max(float(np.linalg.norm(np.ptp(centers, axis=0))), 0.1)
    depth, radius = span * 0.045, span * 0.0015

    def line(a, b, color):
        if np.linalg.norm(b - a) < 1e-9:
            return
        mesh = trimesh.creation.cylinder(radius=radius, segment=np.stack([a, b]), sections=6)
        mesh.visual.face_colors = color
        scene.add_geometry(mesh)

    for index, (pose, k) in enumerate(zip(c2w, intrinsics)):
        pixels = np.array([[0, 0, 1], [width, 0, 1], [width, height, 1], [0, height, 1]])
        corners = (np.linalg.inv(k) @ pixels.T).T * depth
        corners = corners @ pose[:3, :3].T + pose[:3, 3]
        color = [55, int(115 + 105 * index / max(len(c2w) - 1, 1)), 245, 255]
        for corner in corners:
            line(pose[:3, 3], corner, color)
        for left, right in zip(corners, np.roll(corners, -1, axis=0)):
            line(left, right, color)
    if trajectory is not None:
        for left, right in zip(trajectory[:-1, :3, 3], trajectory[1:, :3, 3]):
            line(left, right, [255, 165, 45, 255])
    # Display camera coordinates in a Y-up viewer; exported poses stay untouched.
    scene.apply_transform(np.diag([1.0, -1.0, -1.0, 1.0]))
    output = root / ("cameras.glb" if trajectory is None else "cameras_with_path.glb")
    scene.export(str(output))
    return str(output)