#!/usr/bin/env python3
"""Fuse accepted recorded image masks into calibrated K1 LiDAR observations."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from uuid import uuid4

import numpy as np
import numpy.typing as npt
import rerun as rr
from PIL import Image, ImageDraw
from rerun.components import FillMode

from k1link.device_plugins.xgrids_k1.analyze.calibrated_overlay import (
    _extract_camera_frames,
    _load_calibration_snapshot,
    _select_camera_anchors,
    _select_lidar_samples,
)
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.sessions import inspect_recorded_media_epoch

SCHEMA = "missioncore.calibrated-segmentation-fusion-experiment/v1"
ACCEPTED_INSTANCE_MODEL = "torchvision/maskrcnn_resnet50_fpn_v2"
ACCEPTED_SEMANTIC_MODEL = "microsoft/beit-base-finetuned-ade-640-640"
MIN_BOX_POINTS = 4
MAX_BOX_SPAN_M = 12.0

FloatArray = npt.NDArray[np.float64]
RgbImage = npt.NDArray[np.uint8]


@dataclass(frozen=True, slots=True)
class ObjectFusion:
    instance_id: int
    label: str
    score: float
    candidate_points: int
    clustered_points: int
    distance_p10_m: float | None
    distance_median_m: float | None
    box_center_map: tuple[float, float, float] | None
    box_half_size_map: tuple[float, float, float] | None
    box_status: str
    source_indices: npt.NDArray[np.int64]


@dataclass(frozen=True, slots=True)
class FusedFrame:
    stem: str
    session_time_seconds: float
    image_rgb: RgbImage
    semantic_overlay_rgb: RgbImage
    fusion_overlay_rgb: RgbImage
    points_map: FloatArray
    projected_source_indices: npt.NDArray[np.int64]
    projected_pixels: FloatArray
    semantic_colors: RgbImage
    objects: tuple[ObjectFusion, ...]


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--session", type=Path, required=True)
    parser.add_argument("--calibration", type=Path, required=True)
    parser.add_argument("--calibrated-manifest", type=Path, required=True)
    parser.add_argument("--segmentation", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--ffmpeg", type=Path, required=True)
    return parser.parse_args()


def _read_json(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise RuntimeError(f"{path.name} is not a JSON object")
    return value


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _palette(index: int) -> tuple[int, int, int]:
    digest = hashlib.sha256(f"mission-core-segment-{index}".encode()).digest()
    return (64 + digest[0] % 176, 64 + digest[1] % 176, 64 + digest[2] % 176)


def _depth_cluster(indices: npt.NDArray[np.int64], depths: FloatArray) -> npt.NDArray[np.int64]:
    if indices.size < 2:
        return indices
    order = np.argsort(depths[indices])
    ordered = indices[order]
    ordered_depths = depths[ordered]
    groups: list[npt.NDArray[np.int64]] = []
    start = 0
    for offset, gap in enumerate(np.diff(ordered_depths), start=1):
        threshold = max(0.6, 0.08 * float(ordered_depths[offset - 1]))
        if float(gap) > threshold:
            groups.append(ordered[start:offset])
            start = offset
    groups.append(ordered[start:])
    return min(
        groups,
        key=lambda group: (-int(group.size), float(np.median(depths[group]))),
    )


def _object_fusions(
    *,
    instance_map: npt.NDArray[np.uint16],
    instances: list[dict[str, Any]],
    pixel_x: npt.NDArray[np.int64],
    pixel_y: npt.NDArray[np.int64],
    depths: FloatArray,
    projected_source_indices: npt.NDArray[np.int64],
    points_map: FloatArray,
    points_lidar: FloatArray,
) -> tuple[ObjectFusion, ...]:
    sampled_instances = instance_map[pixel_y, pixel_x]
    fused: list[ObjectFusion] = []
    for item in instances:
        instance_id = int(item["instance_id"])
        candidates = np.flatnonzero(sampled_instances == instance_id).astype(np.int64)
        clustered = _depth_cluster(candidates, depths)
        source_indices = projected_source_indices[clustered]
        ranges = np.linalg.norm(points_lidar[source_indices], axis=1)
        distance_p10 = None if ranges.size == 0 else float(np.percentile(ranges, 10.0))
        distance_median = None if ranges.size == 0 else float(np.median(ranges))
        center: tuple[float, float, float] | None = None
        half_size: tuple[float, float, float] | None = None
        if source_indices.size < MIN_BOX_POINTS:
            status = f"rejected-fewer-than-{MIN_BOX_POINTS}-clustered-points"
        else:
            selected = points_map[source_indices]
            lower = np.percentile(selected, 5.0, axis=0)
            upper = np.percentile(selected, 95.0, axis=0)
            spans = upper - lower
            if not np.isfinite(spans).all() or np.any(spans > MAX_BOX_SPAN_M):
                status = "rejected-implausible-span"
            else:
                sizes = np.maximum(spans, 0.15)
                center = tuple(float(value) for value in ((lower + upper) * 0.5))
                half_size = tuple(float(value) for value in (sizes * 0.5))
                status = "accepted-diagnostic-axis-aligned"
        fused.append(
            ObjectFusion(
                instance_id=instance_id,
                label=str(item["label"]),
                score=float(item["score"]),
                candidate_points=int(candidates.size),
                clustered_points=int(clustered.size),
                distance_p10_m=distance_p10,
                distance_median_m=distance_median,
                box_center_map=center,
                box_half_size_map=half_size,
                box_status=status,
                source_indices=source_indices,
            )
        )
    return tuple(fused)


def _fusion_overlay(
    semantic_overlay: RgbImage,
    projected_pixels: FloatArray,
    semantic_colors: RgbImage,
    objects: tuple[ObjectFusion, ...],
    instances: list[dict[str, Any]],
) -> RgbImage:
    canvas = Image.fromarray(semantic_overlay).convert("RGBA")
    draw = ImageDraw.Draw(canvas, "RGBA")
    for (u, v), color in zip(projected_pixels, semantic_colors, strict=True):
        red, green, blue = (int(value) for value in color)
        draw.ellipse((u - 1.5, v - 1.5, u + 1.5, v + 1.5), fill=(red, green, blue, 210))
    by_id = {item.instance_id: item for item in objects}
    for item in instances:
        fused = by_id[int(item["instance_id"])]
        x1, y1, x2, y2 = (float(value) for value in item["box_xyxy"])
        color = _palette(500 + fused.instance_id)
        accepted = fused.box_center_map is not None
        alpha = 255 if accepted else 135
        draw.rectangle((x1, y1, x2, y2), outline=(*color, alpha), width=2)
        distance = (
            "no LiDAR"
            if fused.distance_median_m is None
            else f"{fused.distance_median_m:.1f}m/{fused.clustered_points}pts"
        )
        draw.text(
            (x1 + 2, max(0.0, y1 - 12)),
            f"{fused.label} {distance}",
            fill=(255, 255, 255, alpha),
            stroke_width=2,
            stroke_fill=(0, 0, 0, alpha),
        )
    return np.asarray(canvas.convert("RGB"), dtype=np.uint8)


def _object_document(item: ObjectFusion) -> dict[str, Any]:
    return {
        "instance_id": item.instance_id,
        "label": item.label,
        "score": round(item.score, 9),
        "candidate_projected_points": item.candidate_points,
        "clustered_points": item.clustered_points,
        "distance_p10_m": item.distance_p10_m,
        "distance_median_m": item.distance_median_m,
        "box_status": item.box_status,
        "box_center_map": item.box_center_map,
        "box_half_size_map": item.box_half_size_map,
    }


def _write_rerun(path: Path, experiment_id: str, frames: tuple[FusedFrame, ...]) -> None:
    recording = rr.RecordingStream("nodedc_k1_calibrated_segmentation", recording_id=experiment_id)
    recording.set_sinks(rr.FileSink(path, write_footer=True))
    try:
        recording.log(
            "contract",
            rr.TextDocument(
                "Recorded-only diagnostic: exact image masks sampled at factory-calibrated "
                "K1 LiDAR projections; 3D boxes require clustered LiDAR support."
            ),
            static=True,
        )
        for frame in frames:
            recording.set_time(
                "session_time",
                duration=np.timedelta64(round(frame.session_time_seconds * 1_000_000_000), "ns"),
            )
            recording.log("camera/raw", rr.Image(frame.image_rgb))
            recording.log("camera/semantic", rr.Image(frame.semantic_overlay_rgb))
            recording.log("camera/fusion", rr.Image(frame.fusion_overlay_rgb))
            recording.log(
                "camera/projected_lidar_semantic",
                rr.Points2D(
                    frame.projected_pixels.astype(np.float32),
                    colors=frame.semantic_colors,
                    radii=rr.Radius.ui_points(2.0),
                ),
            )
            recording.log(
                "world/semantic_points",
                rr.Points3D(
                    frame.points_map[frame.projected_source_indices].astype(np.float32),
                    colors=frame.semantic_colors,
                    radii=rr.Radius.ui_points(2.0),
                ),
            )
            boxed = [item for item in frame.objects if item.box_center_map is not None]
            if boxed:
                recording.log(
                    "world/boxes3d",
                    rr.Boxes3D(
                        centers=[item.box_center_map for item in boxed],
                        half_sizes=[item.box_half_size_map for item in boxed],
                        colors=[(*_palette(500 + item.instance_id), 96) for item in boxed],
                        fill_mode=FillMode.Solid,
                        labels=[
                            f"{item.label} · {item.distance_median_m:.1f} m · "
                            f"{item.clustered_points} pts"
                            for item in boxed
                        ],
                        show_labels=True,
                    ),
                )
    finally:
        recording.flush(timeout_sec=30.0)
        recording.disconnect()


def main() -> int:
    args = _arguments()
    session = args.session.resolve(strict=True)
    calibration_root = args.calibration.resolve(strict=True)
    calibrated_manifest = _read_json(args.calibrated_manifest.resolve(strict=True))
    segmentation_root = args.segmentation.resolve(strict=True)
    segmentation_manifest = _read_json(segmentation_root / "manifest.redacted.json")
    ffmpeg = args.ffmpeg.resolve(strict=True)

    models = segmentation_manifest.get("models")
    acceptance = segmentation_manifest.get("acceptance")
    if not isinstance(models, dict) or not isinstance(acceptance, dict):
        raise RuntimeError("segmentation manifest is incomplete")
    instance_model = models.get("instance")
    semantic_model = models.get("semantic")
    if (
        not isinstance(instance_model, dict)
        or instance_model.get("id") != ACCEPTED_INSTANCE_MODEL
        or not isinstance(semantic_model, dict)
        or semantic_model.get("id") != ACCEPTED_SEMANTIC_MODEL
        or semantic_model.get("checkpoint_load") != "exact-no-missing-unexpected-or-mismatched-keys"
        or acceptance.get("recorded_segmentation_artifact") != "generated"
    ):
        raise RuntimeError("segmentation result is not an admitted strict-load experiment")

    calibration, calibration_identity = _load_calibration_snapshot(calibration_root)
    calibrated_input = calibrated_manifest.get("input")
    if (
        not isinstance(calibrated_input, dict)
        or calibrated_input.get("session_id") != session.name
        or calibrated_input.get("calibration_content_identity_sha256") != calibration_identity
    ):
        raise RuntimeError("calibrated experiment does not bind this session/calibration")
    source_id = str(calibrated_input["source_id"])
    offsets = tuple(float(value) for value in calibrated_input["video_offsets_seconds"])
    profile = Kb4ProjectionProfile.from_factory_calibration(calibration, source_id)

    capture_root = session / "captures" / "mqtt_live"
    origin = read_capture_clock_origin(capture_root / "mqtt.timeline.origin.json")
    epoch_root = session / "media" / source_id / "epoch-1"
    epoch = inspect_recorded_media_epoch(
        epoch_root,
        expected_source_name=source_id,
        origin_epoch_ns=origin.started_at_epoch_ns,
        origin_monotonic_ns=origin.started_monotonic_ns,
    )
    anchors = _select_camera_anchors(
        epoch_root / "index.jsonl",
        expected_count=len(epoch.segments),
        offsets_seconds=offsets,
        timeline_start_seconds=epoch.timeline_start_seconds,
        origin_monotonic_ns=origin.started_monotonic_ns,
    )
    images = _extract_camera_frames(
        ffmpeg,
        init_path=epoch.init_path,
        segment_paths=tuple(item.path for item in epoch.segments),
        frame_sequences=tuple(item.sequence for item in anchors),
        width=profile.width,
        height=profile.height,
    )
    lidar_samples = _select_lidar_samples(
        capture_root / "mqtt.raw.k1mqtt",
        anchors,
        origin_monotonic_ns=origin.started_monotonic_ns,
        temporal_offset_seconds=float(calibrated_input["temporal_offset_seconds"]),
    )

    segmentation_inputs = {
        str(item["name"]): str(item["sha256"])
        for item in segmentation_manifest["input"]
        if isinstance(item, dict)
    }
    segmentation_frames = segmentation_manifest["frames"]
    fused_frames: list[FusedFrame] = []
    for image, anchor, lidar_sample in zip(images, anchors, lidar_samples, strict=True):
        stem = (
            f"camera-{anchor.sequence:06d}-"
            f"session-{round(anchor.session_time_seconds * 1000):09d}ms"
        )
        input_name = f"{stem}.png"
        input_artifact = args.calibrated_manifest.parent / input_name
        if segmentation_inputs.get(input_name) != _sha256(input_artifact):
            raise RuntimeError(f"segmentation input identity changed for {input_name}")
        if _sha256(input_artifact) != hashlib.sha256(image.tobytes()).hexdigest():
            decoded = np.asarray(Image.open(input_artifact).convert("RGB"), dtype=np.uint8)
            if not np.array_equal(decoded, image):
                raise RuntimeError(f"decoded camera pixels changed for {input_name}")

        semantic_map = np.asarray(Image.open(segmentation_root / f"{stem}.semantic.png"))
        instance_map = np.asarray(Image.open(segmentation_root / f"{stem}.instances.png"))
        semantic_overlay = np.asarray(
            Image.open(segmentation_root / f"{stem}.semantic-overlay.png").convert("RGB")
        )
        if semantic_map.shape != image.shape[:2] or instance_map.shape != image.shape[:2]:
            raise RuntimeError("segmentation maps do not match the camera epoch")

        points_map = np.asarray(
            [
                point.scaled_xyz(lidar_sample.point_frame.header.scaler)
                for point in lidar_sample.point_frame.points
            ],
            dtype=np.float64,
        )
        points_lidar = map_points_to_lidar(
            points_map,
            position_map_xyz=lidar_sample.pose_frame.position_xyz,
            orientation_map_from_lidar_xyzw=lidar_sample.pose_frame.orientation_xyzw,
        )
        projection = project_map_points_kb4(
            points_map,
            position_map_xyz=lidar_sample.pose_frame.position_xyz,
            orientation_map_from_lidar_xyzw=lidar_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)
        semantic_colors = np.asarray(
            [_palette(int(value)) for value in semantic_ids], dtype=np.uint8
        )
        frame_document = segmentation_frames[stem]
        instances = frame_document["instance"]["instances"]
        objects = _object_fusions(
            instance_map=instance_map.astype(np.uint16),
            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,
        )
        fusion_overlay = _fusion_overlay(
            semantic_overlay,
            projection.pixels_xy,
            semantic_colors,
            objects,
            instances,
        )
        fused_frames.append(
            FusedFrame(
                stem=stem,
                session_time_seconds=anchor.session_time_seconds,
                image_rgb=image,
                semantic_overlay_rgb=semantic_overlay,
                fusion_overlay_rgb=fusion_overlay,
                points_map=points_map,
                projected_source_indices=projection.source_indices,
                projected_pixels=projection.pixels_xy,
                semantic_colors=semantic_colors,
                objects=objects,
            )
        )

    created_at = datetime.now(UTC)
    experiment_id = (
        f"{created_at.strftime('%Y%m%dT%H%M%SZ')}_k1_segmentation_fusion_{uuid4().hex[:12]}"
    )
    private_root = args.output_root.resolve() / "private" / "perception-experiments"
    private_root.mkdir(mode=0o700, parents=True, exist_ok=True)
    output = private_root / experiment_id
    staging = private_root / f".{experiment_id}.incomplete"
    staging.mkdir(mode=0o700, exist_ok=False)
    try:
        for frame in fused_frames:
            Image.fromarray(frame.fusion_overlay_rgb).save(staging / f"{frame.stem}.fusion.png")
        mosaic_rows = [
            np.concatenate(
                (frame.image_rgb, frame.semantic_overlay_rgb, frame.fusion_overlay_rgb),
                axis=1,
            )
            for frame in fused_frames
        ]
        Image.fromarray(np.concatenate(mosaic_rows, axis=0)).save(staging / "fusion-mosaic.png")
        _write_rerun(staging / "fusion.rrd", experiment_id, tuple(fused_frames))
        identity = {
            "calibrated_generation_sha256": calibrated_manifest["generation_sha256"],
            "segmentation_generation_sha256": segmentation_manifest["generation_sha256"],
            "calibration_content_identity_sha256": calibration_identity,
            "session_id": session.name,
            "source_id": source_id,
        }
        outputs = sorted(staging.iterdir())
        manifest = {
            "schema_version": SCHEMA,
            "experiment_id": experiment_id,
            "created_at_utc": created_at.isoformat(timespec="milliseconds").replace("+00:00", "Z"),
            "generation_sha256": hashlib.sha256(
                json.dumps(identity, sort_keys=True, separators=(",", ":")).encode()
            ).hexdigest(),
            "input": identity,
            "fusion": {
                "pixel_sampling": "nearest-calibrated-projection",
                "occlusion_policy": "largest-depth-contiguous-cluster-nearest-tiebreak",
                "distance": "LiDAR-origin Euclidean p10 and median",
                "box": {
                    "frame": "map",
                    "orientation": "axis-aligned-diagnostic",
                    "minimum_clustered_points": MIN_BOX_POINTS,
                    "robust_bounds": "p05-p95",
                    "maximum_span_m": MAX_BOX_SPAN_M,
                },
            },
            "frames": [
                {
                    "stem": frame.stem,
                    "session_time_seconds": frame.session_time_seconds,
                    "projected_points": int(frame.projected_source_indices.size),
                    "semantic_class_ids": sorted(
                        {
                            int(value)
                            for value in np.asarray(
                                Image.open(segmentation_root / f"{frame.stem}.semantic.png")
                            ).reshape(-1)
                        }
                    ),
                    "objects": [_object_document(item) for item in frame.objects],
                    "accepted_boxes": sum(
                        item.box_center_map is not None for item in frame.objects
                    ),
                }
                for frame in fused_frames
            ],
            "acceptance": {
                "recorded_mask_to_lidar_fusion": "generated",
                "distance": "diagnostic-not-ground-truthed",
                "boxes3d": "diagnostic-not-tracked-or-ground-truthed",
                "live": "not-tested",
                "safety": "not-accepted",
            },
            "outputs": [
                {"name": path.name, "bytes": path.stat().st_size, "sha256": _sha256(path)}
                for path in outputs
            ],
        }
        (staging / "manifest.redacted.json").write_text(
            json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        os.rename(staging, output)
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise
    print(json.dumps({"experiment_id": experiment_id, "output": str(output)}))
    return 0


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