#!/usr/bin/env python3
"""Fuse one admitted full camera-perception epoch into compact 3D observations."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import shutil
import tarfile
import time
from collections import deque
from collections.abc import Iterator
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import numpy as np
import numpy.typing as npt
from fuse_calibrated_segmentation import _object_fusions, _palette
from PIL import Image

from k1link.compute import validate_recorded_perception_epoch_result
from k1link.device_plugins.xgrids_k1.analyze.calibrated_overlay import (
    MAX_SYNC_DELTA_SECONDS,
    _load_calibration_snapshot,
)
from k1link.device_plugins.xgrids_k1.analyze.calibrated_projection import (
    Kb4ProjectionProfile,
    map_points_to_lidar,
    project_map_points_kb4,
)
from k1link.device_plugins.xgrids_k1.mqtt.capture import read_capture_clock_origin
from k1link.device_plugins.xgrids_k1.protocol.streams import (
    LioPointCloudFrame,
    LioPoseFrame,
    decode_lio_pcl,
    decode_lio_pose,
)
from k1link.device_plugins.xgrids_k1.viewer.replay import iter_replay_messages

SCHEMA = "missioncore.recorded-calibrated-fusion/v1"
IDENTITY_SCHEMA = "missioncore.recorded-calibrated-fusion-identity/v1"
MAX_MASK_MEMBER_BYTES = 4 * 1024 * 1024
MAX_INDEX_LINE_BYTES = 64 * 1024

FloatArray = npt.NDArray[np.float64]


@dataclass(frozen=True, slots=True)
class CameraAnchor:
    sequence: int
    host_session_seconds: float
    video_session_seconds: float


@dataclass(frozen=True, slots=True)
class LidarSample:
    point_session_seconds: float
    pose_session_seconds: float
    point_frame: LioPointCloudFrame
    pose_frame: LioPoseFrame


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--job", type=Path, required=True)
    parser.add_argument("--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 _canonical_json(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 chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def _read_frame_metadata(path: Path, expected_count: int) -> list[dict[str, Any]]:
    frames: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as stream:
        for expected_index, line in enumerate(stream):
            value = json.loads(line)
            if not isinstance(value, dict) or value.get("frame_index") != expected_index:
                raise RuntimeError("perception frame metadata changed")
            frames.append(value)
    if len(frames) != expected_count:
        raise RuntimeError("perception frame metadata is incomplete")
    return frames


def _extract_masks(archive: Path, destination: Path, frame_count: int) -> None:
    expected = {
        f"{kind}/frame-{sequence:06d}.png"
        for kind in ("instance-masks", "semantic-masks")
        for sequence in range(1, frame_count + 1)
    }
    with tarfile.open(archive, mode="r:gz") as source:
        members = [member for member in source.getmembers() if member.isfile()]
        found = {member.name.replace("\\", "/").removeprefix("./") for member in members}
        if found != expected or len(members) != len(expected):
            raise RuntimeError("panoptic mask archive contents changed")
        for member in members:
            name = member.name.replace("\\", "/").removeprefix("./")
            if (
                member.issym()
                or member.islnk()
                or member.size < 1
                or member.size > MAX_MASK_MEMBER_BYTES
            ):
                raise RuntimeError("panoptic mask archive member is unsafe")
            target = destination.joinpath(*name.split("/"))
            target.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
            payload = source.extractfile(member)
            if payload is None:
                raise RuntimeError("panoptic mask archive member is unavailable")
            with target.open("xb") as sink:
                shutil.copyfileobj(payload, sink, length=1024 * 1024)
            if target.stat().st_size != member.size:
                raise RuntimeError("panoptic mask archive member was truncated")


def _camera_anchors(
    index_path: Path,
    frame_metadata: list[dict[str, Any]],
    *,
    origin_monotonic_ns: int,
) -> list[CameraAnchor]:
    anchors: list[CameraAnchor] = []
    with index_path.open("rb") as stream:
        for sequence, frame in enumerate(frame_metadata, start=1):
            line = stream.readline(MAX_INDEX_LINE_BYTES + 1)
            if not line or len(line) > MAX_INDEX_LINE_BYTES or not line.endswith(b"\n"):
                raise RuntimeError("camera index is incomplete")
            value = json.loads(line)
            monotonic_ns = value.get("host_monotonic_ns") if isinstance(value, dict) else None
            video_time = frame.get("session_seconds")
            if (
                not isinstance(value, dict)
                or value.get("sequence") != sequence
                or not isinstance(monotonic_ns, int)
                or monotonic_ns < origin_monotonic_ns
                or not isinstance(video_time, (int, float))
                or not math.isfinite(float(video_time))
            ):
                raise RuntimeError("camera anchor is invalid")
            anchors.append(
                CameraAnchor(
                    sequence=sequence,
                    host_session_seconds=(monotonic_ns - origin_monotonic_ns) / 1e9,
                    video_session_seconds=float(video_time),
                )
            )
        if stream.read(1):
            raise RuntimeError("camera index has undeclared rows")
    return anchors


def _lidar_samples(
    raw_path: Path,
    anchors: list[CameraAnchor],
    *,
    origin_monotonic_ns: int,
) -> Iterator[LidarSample | None]:
    messages = iter(iter_replay_messages(raw_path))
    pending = next(messages, None)
    points: deque[tuple[float, LioPointCloudFrame]] = deque()
    poses: deque[tuple[float, LioPoseFrame]] = deque()

    def message_time(message: Any) -> float:
        monotonic_ns = message.received_monotonic_ns
        if not isinstance(monotonic_ns, int) or monotonic_ns < origin_monotonic_ns:
            raise RuntimeError("MQTT replay message has no compatible monotonic time")
        return (monotonic_ns - origin_monotonic_ns) / 1e9

    for anchor in anchors:
        target = anchor.host_session_seconds
        while pending is not None and message_time(pending) <= target + 0.5:
            timestamp = message_time(pending)
            if timestamp >= target - 0.5:
                if pending.topic.endswith("/lio_pcl"):
                    points.append((timestamp, decode_lio_pcl(pending.payload)))
                elif pending.topic.endswith("/lio_pose"):
                    poses.append((timestamp, decode_lio_pose(pending.payload)))
            pending = next(messages, None)
        while points and points[0][0] < target - 0.5:
            points.popleft()
        while poses and poses[0][0] < target - 0.5:
            poses.popleft()
        point_candidates = [
            item for item in points if abs(item[0] - target) <= MAX_SYNC_DELTA_SECONDS
        ]
        if not point_candidates:
            yield None
            continue
        point_time, point_frame = min(point_candidates, key=lambda item: abs(item[0] - target))
        pose_candidates = [
            item for item in poses if abs(item[0] - point_time) <= MAX_SYNC_DELTA_SECONDS
        ]
        if not pose_candidates:
            yield None
            continue
        pose_time, pose_frame = min(pose_candidates, key=lambda item: abs(item[0] - point_time))
        yield LidarSample(point_time, pose_time, point_frame, pose_frame)


def _empty_points() -> npt.NDArray[np.float32]:
    return np.empty((0, 3), dtype=np.float32)


def _empty_colors(channels: int) -> npt.NDArray[np.uint8]:
    return np.empty((0, channels), dtype=np.uint8)


def main() -> int:
    args = _arguments()
    started = time.perf_counter()
    result = validate_recorded_perception_epoch_result(args.job, args.result)
    session = args.session.resolve(strict=True)
    if session.name != result.job.session_id:
        raise RuntimeError("fusion session does not match the perception result")
    calibration, calibration_sha256 = _load_calibration_snapshot(
        args.calibration.resolve(strict=True)
    )
    if calibration_sha256 != result.calibration_sha256:
        raise RuntimeError("fusion calibration generation does not match perception")
    profile = Kb4ProjectionProfile.from_factory_calibration(calibration, result.job.source_id)
    frame_metadata = _read_frame_metadata(
        result.artifact("panoptic-frame-metadata").path,
        result.job.segment_count,
    )
    capture_root = session / "captures" / "mqtt_live"
    origin = read_capture_clock_origin(capture_root / "mqtt.timeline.origin.json")
    epoch_root = (
        result.job.job_root
        / "input"
        / "camera"
        / result.job.source_id
        / f"epoch-{result.job.codec_epoch}"
    )
    anchors = _camera_anchors(
        epoch_root / "index.jsonl",
        frame_metadata,
        origin_monotonic_ns=origin.started_monotonic_ns,
    )
    identity = {
        "schema_version": IDENTITY_SCHEMA,
        "job_id": result.job.job_id,
        "input_sha256": result.job.input_sha256,
        "perception_result_id": result.result_id,
        "calibration_sha256": calibration_sha256,
        "camera_slot": result.calibration_slot,
        "configuration": {
            "pipeline": "factory-kb4-mask-to-lidar/v1",
            "temporal_binding": "nearest-host-arrival-best-effort",
            "maximum_sync_delta_seconds": MAX_SYNC_DELTA_SECONDS,
            "semantic_sampling": "nearest-projected-pixel",
            "box_policy": "depth-clustered-map-axis-aligned-p05-p95-diagnostic",
        },
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    fusion_id = f"fusion-{identity_sha256}"
    parent = args.output_root.resolve() / result.job.job_id
    parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    output = parent / fusion_id
    if output.exists():
        print(json.dumps({"fusion_id": fusion_id, "output": str(output), "reused": True}))
        return 0
    staging = parent / f".{fusion_id}.{os.getpid()}.incomplete"
    staging.mkdir(mode=0o700, exist_ok=False)
    masks_root = staging / "masks"
    masks_root.mkdir(mode=0o700)
    try:
        _extract_masks(
            result.artifact("panoptic-mask-archive").path,
            masks_root,
            result.job.segment_count,
        )
        frame_times_ns: list[int] = []
        point_offsets = [0]
        point_arrays: list[npt.NDArray[np.float32]] = []
        point_color_arrays: list[npt.NDArray[np.uint8]] = []
        box_offsets = [0]
        box_centers: list[tuple[float, float, float]] = []
        box_half_sizes: list[tuple[float, float, float]] = []
        box_colors: list[tuple[int, int, int, int]] = []
        box_labels: list[str] = []
        frame_reports: list[dict[str, Any]] = []
        compatible_frames = 0
        accepted_boxes = 0
        for index, (anchor, sample) in enumerate(
            zip(
                anchors,
                _lidar_samples(
                    capture_root / "mqtt.raw.k1mqtt",
                    anchors,
                    origin_monotonic_ns=origin.started_monotonic_ns,
                ),
                strict=True,
            )
        ):
            frame_times_ns.append(round(anchor.video_session_seconds * 1e9))
            if sample is None:
                point_arrays.append(_empty_points())
                point_color_arrays.append(_empty_colors(3))
                point_offsets.append(point_offsets[-1])
                box_offsets.append(box_offsets[-1])
                frame_reports.append(
                    {"frame_index": index, "state": "depth-unavailable", "points": 0, "boxes": 0}
                )
                continue
            compatible_frames += 1
            with Image.open(
                masks_root / "semantic-masks" / f"frame-{index + 1:06d}.png"
            ) as semantic_image:
                semantic_map = np.array(semantic_image, dtype=np.uint8, copy=True)
            with Image.open(
                masks_root / "instance-masks" / f"frame-{index + 1:06d}.png"
            ) as instance_image:
                instance_map = np.array(instance_image, dtype=np.uint16, copy=True)
            if semantic_map.shape != (profile.height, profile.width) or instance_map.shape != (
                profile.height,
                profile.width,
            ):
                raise RuntimeError("panoptic mask dimensions changed")
            points_map = np.asarray(
                [
                    point.scaled_xyz(sample.point_frame.header.scaler)
                    for point in sample.point_frame.points
                ],
                dtype=np.float64,
            )
            points_lidar = map_points_to_lidar(
                points_map,
                position_map_xyz=sample.pose_frame.position_xyz,
                orientation_map_from_lidar_xyzw=sample.pose_frame.orientation_xyzw,
            )
            projection = project_map_points_kb4(
                points_map,
                position_map_xyz=sample.pose_frame.position_xyz,
                orientation_map_from_lidar_xyzw=sample.pose_frame.orientation_xyzw,
                profile=profile,
            )
            pixel_x = np.clip(
                np.rint(projection.pixels_xy[:, 0]).astype(np.int64),
                0,
                profile.width - 1,
            )
            pixel_y = np.clip(
                np.rint(projection.pixels_xy[:, 1]).astype(np.int64),
                0,
                profile.height - 1,
            )
            semantic_ids = semantic_map[pixel_y, pixel_x].astype(np.int64)
            colors = np.asarray([_palette(int(value)) for value in semantic_ids], dtype=np.uint8)
            projected = points_map[projection.source_indices].astype(np.float32)
            point_arrays.append(projected)
            point_color_arrays.append(colors)
            point_offsets.append(point_offsets[-1] + projected.shape[0])
            instances = frame_metadata[index]["instances"]
            objects = _object_fusions(
                instance_map=instance_map,
                instances=instances,
                pixel_x=pixel_x,
                pixel_y=pixel_y,
                depths=projection.depths_m,
                projected_source_indices=projection.source_indices,
                points_map=points_map,
                points_lidar=points_lidar,
            )
            frame_box_count = 0
            for item in objects:
                if item.box_center_map is None or item.box_half_size_map is None:
                    continue
                frame_box_count += 1
                accepted_boxes += 1
                box_centers.append(item.box_center_map)
                box_half_sizes.append(item.box_half_size_map)
                box_colors.append((*_palette(500 + item.instance_id), 96))
                box_labels.append(
                    f"{item.label} · {item.distance_median_m:.1f} m · {item.clustered_points} pts"
                )
            box_offsets.append(box_offsets[-1] + frame_box_count)
            frame_reports.append(
                {
                    "frame_index": index,
                    "state": "fused",
                    "points": int(projected.shape[0]),
                    "boxes": frame_box_count,
                    "lidar_camera_delta_ms": round(
                        (sample.point_session_seconds - anchor.host_session_seconds) * 1000,
                        6,
                    ),
                    "pose_point_delta_ms": round(
                        (sample.pose_session_seconds - sample.point_session_seconds) * 1000,
                        6,
                    ),
                }
            )
            if (index + 1) % 250 == 0 or index + 1 == result.job.segment_count:
                print(
                    json.dumps(
                        {"phase": "fusion", "frames": index + 1, "total": result.job.segment_count}
                    ),
                    flush=True,
                )
        arrays_path = staging / "fusion.npz"
        np.savez_compressed(
            arrays_path,
            frame_times_ns=np.asarray(frame_times_ns, dtype=np.int64),
            point_offsets=np.asarray(point_offsets, dtype=np.int64),
            points=(np.concatenate(point_arrays) if point_arrays else _empty_points()),
            point_colors=(
                np.concatenate(point_color_arrays) if point_color_arrays else _empty_colors(3)
            ),
            box_offsets=np.asarray(box_offsets, dtype=np.int64),
            box_centers=np.asarray(box_centers, dtype=np.float32).reshape((-1, 3)),
            box_half_sizes=np.asarray(box_half_sizes, dtype=np.float32).reshape((-1, 3)),
            box_colors=np.asarray(box_colors, dtype=np.uint8).reshape((-1, 4)),
        )
        labels_path = staging / "box-labels.json"
        labels_path.write_text(
            json.dumps(box_labels, ensure_ascii=False, separators=(",", ":")),
            encoding="utf-8",
        )
        frames_path = staging / "fusion-frames.jsonl"
        frames_path.write_text(
            "".join(json.dumps(item, sort_keys=True) + "\n" for item in frame_reports),
            encoding="utf-8",
        )
        shutil.rmtree(masks_root)
        created_at = datetime.now(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z")
        artifacts = [
            {
                "name": path.name,
                "byte_length": path.stat().st_size,
                "sha256": _sha256(path),
            }
            for path in (arrays_path, labels_path, frames_path)
        ]
        manifest = {
            "schema_version": SCHEMA,
            "fusion_id": fusion_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "created_at_utc": created_at,
            "session_id": result.job.session_id,
            "source_id": result.job.source_id,
            "frame_count": result.job.segment_count,
            "timeline_start_seconds": result.job.timeline_start_seconds,
            "timeline_end_seconds": result.job.timeline_end_seconds,
            "metrics": {
                "frames_fused": compatible_frames,
                "frames_depth_unavailable": result.job.segment_count - compatible_frames,
                "semantic_points": point_offsets[-1],
                "accepted_diagnostic_boxes": accepted_boxes,
                "wall_seconds": round(time.perf_counter() - started, 6),
            },
            "acceptance": {
                "geometry": "factory-calibrated-kb4",
                "timing": "host-arrival-best-effort-not-ground-truthed",
                "distance": "diagnostic-not-ground-truthed",
                "boxes3d": "support-gated-diagnostic-not-tracked",
                "safety": "not-accepted",
            },
            "artifacts": artifacts,
        }
        (staging / "manifest.json").write_text(
            json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        os.replace(staging, output)
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise
    print(json.dumps({"fusion_id": fusion_id, "output": str(output), "reused": False}))
    return 0


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