#!/usr/bin/env python3
"""Build a content-addressed LAB E10 LiDAR/pose replay pack from accepted E6."""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
import shutil
import sys
from datetime import UTC, datetime
from pathlib import Path
from types import ModuleType

import numpy as np
from e10_fusion_runtime import PACK_SCHEMA, canonical_json, sha256

from k1link.compute import validate_tracked_fusion_qualification_result
from k1link.device_plugins.xgrids_k1.analyze.calibrated_overlay import (
    _load_calibration_snapshot,
)
from k1link.device_plugins.xgrids_k1.analyze.calibrated_projection import (
    Kb4ProjectionProfile,
)


def arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--job", type=Path, required=True)
    parser.add_argument("--tracking", type=Path, required=True)
    parser.add_argument("--semantic", type=Path, required=True)
    parser.add_argument("--e6-result", type=Path, required=True)
    parser.add_argument("--session", type=Path, required=True)
    parser.add_argument("--calibration", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


def load_e7_runner() -> ModuleType:
    path = Path(__file__).with_name("run_e7_replay_realtime.py").resolve(strict=True)
    spec = importlib.util.spec_from_file_location("missioncore_e10_e7_adapter", path)
    if spec is None or spec.loader is None:
        raise RuntimeError("LAB E7 adapter cannot be loaded")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def main() -> int:
    args = arguments()
    upstream = validate_tracked_fusion_qualification_result(
        args.job,
        args.tracking,
        args.semantic,
        args.e6_result,
    )
    session = args.session.resolve(strict=True)
    if session.name != upstream.job.session_id:
        raise RuntimeError("LAB E10 session differs from the accepted E6 input")
    calibration, calibration_sha256 = _load_calibration_snapshot(
        args.calibration.resolve(strict=True)
    )
    if calibration_sha256 != upstream.semantic.calibration_sha256:
        raise RuntimeError("LAB E10 calibration differs from the accepted E6 input")
    projection = Kb4ProjectionProfile.from_factory_calibration(
        calibration,
        upstream.job.source_id,
    )
    adapter = load_e7_runner()
    frames = adapter._prepare_frames(
        e6=upstream,
        session=session,
        selection_start=0,
        selection_count=upstream.frame_count,
    )
    if len(frames) != upstream.frame_count:
        raise RuntimeError("LAB E10 E6 replay frame count changed")

    frame_indices = np.arange(len(frames), dtype=np.int64)
    source_indices = np.asarray(
        [int(frame.frame["source_frame_index"]) for frame in frames], dtype=np.int64
    )
    session_seconds = np.asarray(
        [float(frame.frame["session_seconds"]) for frame in frames], dtype=np.float64
    )
    available = np.asarray([frame.position_map_xyz is not None for frame in frames], dtype=np.bool_)
    offsets = [0]
    clouds = []
    positions = np.full((len(frames), 3), np.nan, dtype=np.float64)
    quaternions = np.full((len(frames), 4), np.nan, dtype=np.float64)
    lidar_delta = np.full((len(frames),), np.nan, dtype=np.float64)
    pose_delta = np.full((len(frames),), np.nan, dtype=np.float64)
    for index, frame in enumerate(frames):
        cloud = np.asarray(frame.cloud_map, dtype=np.float32)
        clouds.append(cloud)
        offsets.append(offsets[-1] + cloud.shape[0])
        if frame.position_map_xyz is not None and frame.orientation_map_from_lidar_xyzw is not None:
            positions[index] = frame.position_map_xyz
            quaternions[index] = frame.orientation_map_from_lidar_xyzw
            lidar_delta[index] = float(frame.frame["lidar_camera_delta_ms"])
            pose_delta[index] = float(frame.frame["pose_point_delta_ms"])

    identity = {
        "schema_version": PACK_SCHEMA,
        "job_id": upstream.job.job_id,
        "input_sha256": upstream.job.input_sha256,
        "session_id": upstream.job.session_id,
        "source_id": upstream.job.source_id,
        "camera_slot": upstream.semantic.calibration_slot,
        "calibration_sha256": calibration_sha256,
        "e6_result_id": upstream.result_id,
        "frame_count": len(frames),
        "source_start_frame_index": int(source_indices[0]),
        "source_end_frame_index": int(source_indices[-1]),
        "timeline_start_seconds": float(session_seconds[0]),
        "timeline_end_seconds": float(session_seconds[-1]),
        "available_lidar_frames": int(available.sum()),
        "point_count": int(offsets[-1]),
        "temporal_binding": "accepted-e6-nearest-host-arrival-best-effort",
        "projection": {
            "model": "kb4",
            "width": projection.width,
            "height": projection.height,
            "source_coordinates": "k1-map",
            "target_camera": upstream.job.source_id,
        },
        "producer_sha256": sha256(Path(__file__).resolve(strict=True)),
    }
    identity_sha256 = hashlib.sha256(canonical_json(identity)).hexdigest()
    pack_id = f"e10-lidar-pack-{identity_sha256}"
    parent = args.output_root.resolve()
    parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    output = parent / pack_id
    if output.exists():
        print(json.dumps({"pack_id": pack_id, "output": str(output), "reused": True}))
        return 0
    staging = parent / f".{pack_id}.{os.getpid()}.incomplete"
    staging.mkdir(mode=0o700, exist_ok=False)
    try:
        arrays_path = staging / "lidar-pack.npz"
        np.savez_compressed(
            arrays_path,
            frame_indices=frame_indices,
            source_frame_indices=source_indices,
            session_seconds=session_seconds,
            sample_available=available,
            cloud_offsets=np.asarray(offsets, dtype=np.int64),
            cloud_points_map=(
                np.concatenate(clouds) if clouds else np.empty((0, 3), dtype=np.float32)
            ),
            pose_positions_map=positions,
            pose_quaternions_map_from_lidar=quaternions,
            lidar_camera_delta_ms=lidar_delta,
            pose_point_delta_ms=pose_delta,
            intrinsic_fx_fy_cx_cy=np.asarray(projection.intrinsic_fx_fy_cx_cy, dtype=np.float64),
            distortion_kb4=np.asarray(projection.distortion_kb4, dtype=np.float64),
            t_camera_from_lidar=np.asarray(projection.t_camera_from_lidar, dtype=np.float64),
        )
        manifest = {
            "schema_version": PACK_SCHEMA,
            "pack_id": pack_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "created_at_utc": datetime.now(UTC)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z"),
            "classification": "private-recorded-sensor-replay-input",
            "ground_truth": False,
            "artifact": {
                "path": arrays_path.name,
                "media_type": "application/x-npz",
                "byte_length": arrays_path.stat().st_size,
                "sha256": sha256(arrays_path),
            },
        }
        manifest_path = staging / "manifest.json"
        manifest_path.write_bytes(
            json.dumps(
                manifest,
                ensure_ascii=False,
                sort_keys=True,
                indent=2,
                allow_nan=False,
            ).encode()
            + b"\n"
        )
        os.replace(staging, output)
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise
    print(
        json.dumps(
            {
                "pack_id": pack_id,
                "output": str(output),
                "frames": len(frames),
                "lidar_frames": int(available.sum()),
                "points": int(offsets[-1]),
                "byte_length": (output / "lidar-pack.npz").stat().st_size,
            },
            sort_keys=True,
        )
    )
    return 0


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