#!/usr/bin/env python3
"""Build a full-camera-epoch LAB E10 LiDAR/pose replay pack from raw evidence."""

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
from e10_fusion_runtime import PACK_SCHEMA, LidarReplayPack, canonical_json, sha256
from fuse_e6_tracking_lidar import _camera_anchors, _lidar_samples, _read_profile

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,
)
from k1link.device_plugins.xgrids_k1.mqtt.capture import read_capture_clock_origin


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("--e6-profile", 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 _read_e6_identity(result_root: Path) -> dict[str, Any]:
    document = json.loads(
        (result_root.resolve(strict=True) / "result.json").read_text(encoding="utf-8")
    )
    identity = document.get("identity") if isinstance(document, dict) else None
    if not isinstance(identity, dict):
        raise RuntimeError("accepted LAB E6 identity is unavailable")
    return identity


def _read_full_timeline(path: Path, expected_count: int) -> list[dict[str, Any]]:
    frames: list[dict[str, Any]] = []
    previous_seconds = -1.0
    with path.resolve(strict=True).open(encoding="utf-8") as stream:
        for frame_index, line in enumerate(stream):
            value = json.loads(line)
            session_seconds = value.get("session_seconds") if isinstance(value, dict) else None
            if (
                not isinstance(value, dict)
                or value.get("frame_index") != frame_index
                or value.get("sequence") != frame_index + 1
                or not isinstance(session_seconds, (int, float))
                or isinstance(session_seconds, bool)
                or not float(session_seconds) > previous_seconds
            ):
                raise RuntimeError("full camera timeline is invalid")
            frames.append(
                {
                    "frame_index": frame_index,
                    "source_frame_index": frame_index,
                    "session_seconds": float(session_seconds),
                    "tracks": [],
                }
            )
            previous_seconds = float(session_seconds)
    if len(frames) != expected_count:
        raise RuntimeError("full camera timeline count changed")
    return frames


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 full session differs from accepted E6")
    e6_profile, e6_profile_sha256 = _read_profile(args.e6_profile)
    e6_identity = _read_e6_identity(upstream.result_root)
    if (
        e6_identity.get("profile_sha256") != e6_profile_sha256
        or e6_identity.get("configuration") != e6_profile
    ):
        raise RuntimeError("LAB E10 full pack temporal policy differs from accepted E6")
    calibration, calibration_sha256 = _load_calibration_snapshot(
        args.calibration.resolve(strict=True)
    )
    if calibration_sha256 != upstream.semantic.calibration_sha256:
        raise RuntimeError("LAB E10 full pack calibration differs from accepted E6")
    projection = Kb4ProjectionProfile.from_factory_calibration(
        calibration,
        upstream.job.source_id,
    )
    timeline_frames = _read_full_timeline(
        upstream.semantic.artifact("panoptic-frame-metadata").path,
        upstream.job.segment_count,
    )
    capture_root = session / "captures" / "mqtt_live"
    origin = read_capture_clock_origin(capture_root / "mqtt.timeline.origin.json")
    epoch_root = (
        upstream.job.job_root
        / "input"
        / "camera"
        / upstream.job.source_id
        / f"epoch-{upstream.job.codec_epoch}"
    )
    anchors = _camera_anchors(
        epoch_root / "index.jsonl",
        timeline_frames,
        origin_monotonic_ns=origin.started_monotonic_ns,
    )
    maximum_lidar_delta_s = float(e6_profile["temporal"]["maximum_lidar_camera_delta_ms"]) / 1000.0
    maximum_pose_delta_s = float(e6_profile["temporal"]["maximum_pose_point_delta_ms"]) / 1000.0
    samples = _lidar_samples(
        capture_root / "mqtt.raw.k1mqtt",
        anchors,
        origin_monotonic_ns=origin.started_monotonic_ns,
        maximum_lidar_camera_delta_s=maximum_lidar_delta_s,
        maximum_pose_point_delta_s=maximum_pose_delta_s,
    )

    count = len(anchors)
    available = np.zeros((count,), dtype=np.bool_)
    offsets = [0]
    clouds: list[np.ndarray] = []
    positions = np.full((count, 3), np.nan, dtype=np.float64)
    quaternions = np.full((count, 4), np.nan, dtype=np.float64)
    lidar_delta = np.full((count,), np.nan, dtype=np.float64)
    pose_delta = np.full((count,), np.nan, dtype=np.float64)
    for index, (anchor, sample) in enumerate(zip(anchors, samples, strict=True)):
        if sample is None:
            offsets.append(offsets[-1])
        else:
            cloud = np.asarray(
                [
                    point.scaled_xyz(sample.point_frame.header.scaler)
                    for point in sample.point_frame.points
                ],
                dtype=np.float32,
            ).reshape((-1, 3))
            available[index] = True
            clouds.append(cloud)
            offsets.append(offsets[-1] + cloud.shape[0])
            positions[index] = sample.pose_frame.position_xyz
            quaternions[index] = sample.pose_frame.orientation_xyzw
            lidar_delta[index] = (
                sample.point_session_seconds - anchor.host_session_seconds
            ) * 1000.0
            pose_delta[index] = (
                sample.pose_session_seconds - sample.point_session_seconds
            ) * 1000.0
        if (index + 1) % 500 == 0 or index + 1 == count:
            print(
                json.dumps(
                    {
                        "phase": "full-lidar-pack",
                        "frames": index + 1,
                        "total": count,
                        "lidar_frames": int(available[: index + 1].sum()),
                        "points": offsets[-1],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )

    source_indices = np.arange(count, dtype=np.int64)
    session_seconds = np.asarray(
        [anchor.video_session_seconds for anchor in anchors], dtype=np.float64
    )
    temporal = {
        "binding": e6_profile["temporal"]["binding"],
        "maximum_lidar_camera_delta_ms": float(
            e6_profile["temporal"]["maximum_lidar_camera_delta_ms"]
        ),
        "maximum_pose_point_delta_ms": float(e6_profile["temporal"]["maximum_pose_point_delta_ms"]),
        "clock_source": "recorded-host-monotonic-arrival",
    }
    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,
        "e6_profile_sha256": e6_profile_sha256,
        "semantic_timeline_result_id": upstream.semantic.result_id,
        "frame_count": count,
        "source_start_frame_index": 0,
        "source_end_frame_index": count - 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",
        "temporal_policy": temporal,
        "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=np.arange(count, dtype=np.int64),
            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),
            },
        }
        (staging / "manifest.json").write_bytes(
            json.dumps(
                manifest,
                ensure_ascii=False,
                sort_keys=True,
                indent=2,
                allow_nan=False,
            ).encode()
            + b"\n"
        )
        os.replace(staging, output)
        try:
            validation = LidarReplayPack(output, expected_job_id=upstream.job.job_id)
            validation.close()
        except BaseException:
            shutil.rmtree(output, ignore_errors=True)
            raise
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise
    print(
        json.dumps(
            {
                "pack_id": pack_id,
                "output": str(output),
                "frames": count,
                "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())
