#!/usr/bin/env python3
"""Extract a minimal content-addressed live projection pack from accepted E14 input."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import numpy as np

SCHEMA = "missioncore.e15-live-projection-pack/v1"
PACK_ARRAYS = {
    "intrinsic_fx_fy_cx_cy",
    "distortion_kb4",
    "t_camera_from_lidar",
}


def _canonical(value: object) -> bytes:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode()


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while block := stream.read(1024 * 1024):
            digest.update(block)
    return digest.hexdigest()


def _read_object(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise RuntimeError("projection source manifest is invalid")
    return value


def build_projection_pack(
    source_root: Path,
    output_root: Path,
    *,
    expected_calibration_sha256: str,
) -> Path:
    source = source_root.expanduser().resolve(strict=True)
    manifest_path = source / "manifest.json"
    arrays_path = source / "lidar-pack.npz"
    manifest = _read_object(manifest_path)
    identity = manifest.get("identity")
    if (
        not isinstance(identity, dict)
        or identity.get("calibration_sha256") != expected_calibration_sha256
        or identity.get("source_id") != "sensor.camera.right"
        or identity.get("camera_slot") != "camera_1"
        or manifest.get("artifact", {}).get("sha256") != _sha256(arrays_path)
    ):
        raise RuntimeError("accepted LiDAR pack calibration binding changed")
    with np.load(arrays_path, allow_pickle=False) as source_arrays:
        if not PACK_ARRAYS.issubset(source_arrays.files):
            raise RuntimeError("accepted LiDAR pack projection arrays changed")
        intrinsic = np.asarray(source_arrays["intrinsic_fx_fy_cx_cy"], dtype=np.float64)
        distortion = np.asarray(source_arrays["distortion_kb4"], dtype=np.float64)
        transform = np.asarray(source_arrays["t_camera_from_lidar"], dtype=np.float64)
    if (
        intrinsic.shape != (4,)
        or distortion.shape != (4,)
        or transform.shape != (4, 4)
        or not np.isfinite(intrinsic).all()
        or not np.isfinite(distortion).all()
        or not np.isfinite(transform).all()
    ):
        raise RuntimeError("live projection arrays are invalid")

    pack_identity = {
        "schema_version": SCHEMA,
        "source_id": "sensor.camera.right",
        "calibration_slot": "camera_1",
        "calibration_sha256": expected_calibration_sha256,
        "resolution": [800, 600],
        "projection_model": "kb4",
        "coordinate_contract": "k1-map-to-camera_1-via-inverse-map-pose",
        "source_lidar_pack_id": manifest.get("pack_id"),
        "source_lidar_pack_identity_sha256": manifest.get("identity_sha256"),
        "producer_sha256": _sha256(Path(__file__).resolve(strict=True)),
    }
    identity_sha256 = hashlib.sha256(_canonical(pack_identity)).hexdigest()
    pack_id = f"e15-live-projection-{identity_sha256}"
    parent = output_root.expanduser().resolve()
    parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    destination = parent / pack_id
    if destination.exists():
        validate_projection_pack(destination, expected_calibration_sha256)
        return destination
    staging = parent / f".{pack_id}.{os.getpid()}.incomplete"
    staging.mkdir(mode=0o700, exist_ok=False)
    try:
        output_arrays = staging / "projection.npz"
        np.savez(
            output_arrays,
            intrinsic_fx_fy_cx_cy=intrinsic,
            distortion_kb4=distortion,
            t_camera_from_lidar=transform,
        )
        document = {
            "schema_version": SCHEMA,
            "pack_id": pack_id,
            "identity_sha256": identity_sha256,
            "identity": pack_identity,
            "created_at_utc": datetime.now(UTC)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z"),
            "classification": "calibration-only-no-recorded-sensor-frames",
            "artifact": {
                "path": output_arrays.name,
                "byte_length": output_arrays.stat().st_size,
                "sha256": _sha256(output_arrays),
            },
        }
        (staging / "manifest.json").write_bytes(
            json.dumps(document, indent=2, sort_keys=True).encode() + b"\n"
        )
        os.replace(staging, destination)
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise
    validate_projection_pack(destination, expected_calibration_sha256)
    return destination


def validate_projection_pack(root: Path, expected_calibration_sha256: str) -> dict[str, Any]:
    resolved = root.expanduser().resolve(strict=True)
    manifest = _read_object(resolved / "manifest.json")
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    artifact = manifest.get("artifact")
    arrays_path = resolved / "projection.npz"
    if (
        manifest.get("schema_version") != SCHEMA
        or not isinstance(identity, dict)
        or identity.get("schema_version") != SCHEMA
        or identity.get("calibration_sha256") != expected_calibration_sha256
        or not isinstance(identity_sha256, str)
        or hashlib.sha256(_canonical(identity)).hexdigest() != identity_sha256
        or resolved.name != f"e15-live-projection-{identity_sha256}"
        or manifest.get("pack_id") != resolved.name
        or not isinstance(artifact, dict)
        or artifact.get("path") != arrays_path.name
        or artifact.get("byte_length") != arrays_path.stat().st_size
        or artifact.get("sha256") != _sha256(arrays_path)
    ):
        raise RuntimeError("live projection pack identity is invalid")
    with np.load(arrays_path, allow_pickle=False) as arrays:
        if set(arrays.files) != PACK_ARRAYS:
            raise RuntimeError("live projection pack arrays changed")
        if (
            arrays["intrinsic_fx_fy_cx_cy"].shape != (4,)
            or arrays["distortion_kb4"].shape != (4,)
            or arrays["t_camera_from_lidar"].shape != (4, 4)
        ):
            raise RuntimeError("live projection pack shapes changed")
    return manifest


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-root", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--calibration-sha256", required=True)
    return parser.parse_args()


def main() -> int:
    args = _arguments()
    output = build_projection_pack(
        args.source_root,
        args.output_root,
        expected_calibration_sha256=args.calibration_sha256,
    )
    print(json.dumps({"projection_pack": str(output)}, sort_keys=True))
    return 0


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