from __future__ import annotations

import errno
import logging
import math
import socket
import time
from collections.abc import Callable, Mapping
from contextlib import suppress
from dataclasses import asdict, dataclass
from typing import Any, Literal
from uuid import uuid4

import numpy as np
import rerun as rr
from rerun import blueprint as rrb
from rerun.components import FillMode

from k1link.compute.live_perception import LivePerceptionResultFrame
from k1link.data_plane import (
    DecodedDataPlaneView,
    DecodedPointCloudView,
    DecodedPoseView,
)
from k1link.viewer.metrics import BridgeMetrics

PointColorMode = Literal["intensity", "height", "distance", "rgb", "class"]
PointPalette = Literal["turbo", "viridis", "plasma", "grayscale", "custom"]

MAX_TRAJECTORY_POSES = 2_000
TRAJECTORY_APPEND_INTERVAL_NS = 500_000_000
TRAJECTORY_FORCE_APPEND_NS = 2_000_000_000
TRAJECTORY_MIN_DISTANCE_METERS = 0.02
TRAJECTORY_PUBLISH_INTERVAL_NS = 500_000_000
LIVE_GRPC_BUFFER_LIMIT = "32MiB"
DEFAULT_GRPC_PORT = 9876
GRPC_PORT_SEARCH_SPAN = 128
DEFAULT_CORS_ORIGINS = (
    "http://127.0.0.1:5173",
    "http://localhost:5173",
    "http://127.0.0.1:4173",
    "http://localhost:4173",
    "http://127.0.0.1:8000",
    "http://localhost:8000",
)

logger = logging.getLogger("k1link.viewer.rerun_bridge")


@dataclass(frozen=True, slots=True)
class RerunSceneSettings:
    point_size: float = 2.5
    color_mode: PointColorMode = "intensity"
    palette: PointPalette = "turbo"
    custom_color: str = "#f7f8f4"
    accumulation_seconds: float = 12.0
    show_points: bool = True
    show_trajectory: bool = True
    show_grid: bool = True
    show_detections_2d: bool = False
    show_segmentation: bool = False
    show_cuboids_3d: bool = False

    def as_dict(self) -> dict[str, object]:
        return asdict(self)


SettingsProvider = Callable[[], RerunSceneSettings]


def _select_available_grpc_port(
    preferred_port: int,
    *,
    search_span: int = GRPC_PORT_SEARCH_SPAN,
) -> int:
    """Select a local Rerun port without reusing a still-served recording."""

    if not 1 <= preferred_port <= 65_535:
        raise ValueError("Rerun gRPC port must be between 1 and 65535")
    if search_span < 1:
        raise ValueError("Rerun gRPC port search span must be positive")

    last_port = min(preferred_port + search_span, 65_536)
    for candidate in range(preferred_port, last_port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
            try:
                # Rerun binds all local interfaces. Probe the same address class
                # so a previous recording retained by a late viewer is detected.
                probe.bind(("0.0.0.0", candidate))
            except OSError as exc:
                if exc.errno == errno.EADDRINUSE:
                    continue
                if exc.errno in {errno.EACCES, errno.EPERM}:
                    raise PermissionError(
                        exc.errno,
                        "Permission denied while probing local Rerun gRPC "
                        f"port {candidate}",
                    ) from exc
                raise RuntimeError(
                    "Could not probe local Rerun gRPC "
                    f"port {candidate}: {type(exc).__name__}: {exc}"
                ) from exc
        return candidate
    raise RuntimeError(
        "No local Rerun gRPC port is available in "
        f"{preferred_port}..{last_port - 1}"
    )


class RerunBridge:
    """Publish transport-neutral canonical envelopes to a Rerun recording."""

    def __init__(
        self,
        *,
        grpc_port: int = DEFAULT_GRPC_PORT,
        metrics: BridgeMetrics | None = None,
        settings_provider: SettingsProvider | None = None,
        cors_allow_origin: tuple[str, ...] = DEFAULT_CORS_ORIGINS,
        recording_factory: Callable[[str], rr.RecordingStream] | None = None,
    ) -> None:
        self.metrics = metrics or BridgeMetrics()
        self._settings_provider = settings_provider or RerunSceneSettings
        self._settings = self._settings_provider()
        if recording_factory is None:
            recording = rr.RecordingStream(
                "nodedc_mission_core_spatial",
                recording_id=uuid4(),
            )
        else:
            recording = recording_factory("nodedc_mission_core_spatial")
        try:
            selected_grpc_port = _select_available_grpc_port(grpc_port)
            if selected_grpc_port != grpc_port:
                logger.info(
                    "Mission Core selected a new Rerun port because an earlier "
                    "viewer still owns the preferred listener",
                    extra={
                        "event_code": "rerun_grpc_port_rotated",
                        "preferred_port": grpc_port,
                        "selected_port": selected_grpc_port,
                    },
                )
            blueprint = _blueprint(self._settings)
            url = recording.serve_grpc(
                grpc_port=selected_grpc_port,
                default_blueprint=blueprint,
                # This is a reconnect cushion for the live preview, not the source
                # of record. Raw MQTT evidence is persisted independently. A large
                # late-client backlog can block the native SDK and freeze preview.
                server_memory_limit=LIVE_GRPC_BUFFER_LIMIT,
                # Rerun 0.34.1 can replay ActivateStore before StoreInfo when an
                # evicted buffer is served newest-first, leaving late viewers on the
                # welcome screen. Preserve protocol order within the bounded cache.
                newest_first=False,
                cors_allow_origin=list(cors_allow_origin),
            )
            recording.send_blueprint(
                blueprint,
                make_active=True,
                make_default=True,
            )
            recording.log("/world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
            recording.log(
                "/world/sensor_pose",
                rr.TransformAxes3D(axis_length=0.45, show_frame=False),
                static=True,
            )
            # Do not expose a URL whose StoreInfo, blueprint and static scene are
            # still waiting in the SDK micro-batcher.
            recording.flush(timeout_sec=5.0)
        except BaseException:
            # Preserve the construction failure while still making a best
            # effort to release a partially started native listener.
            with suppress(BaseException):
                recording.disconnect()
            raise
        self._recording = recording
        self._url = url
        self._path: list[tuple[float, float, float]] = []
        self._last_path_append_source_ns = 0
        self._last_trajectory_publish_ns = 0
        self._last_point_count = 0
        self._closed = False

    @property
    def grpc_url(self) -> str:
        return self._url

    def begin_session(self, metrics: BridgeMetrics | None = None) -> None:
        """Initialize the one acquisition owned by this recording server."""
        if self._closed:
            raise RuntimeError("Rerun bridge is already closed")
        if metrics is not None:
            self.metrics = metrics
        self._settings = self._settings_provider()
        self._path.clear()
        self._last_path_append_source_ns = 0
        self._last_trajectory_publish_ns = 0
        self._last_point_count = 0
        self._recording.set_time("stream_time", timestamp=time.time())
        self._recording.log("/world", rr.Clear(recursive=True))
        self._recording.log("/world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
        self._recording.log(
            "/world/sensor_pose",
            rr.TransformAxes3D(axis_length=0.45, show_frame=False),
            static=True,
        )
        self._recording.send_blueprint(
            _blueprint(self._settings),
            make_active=True,
            make_default=True,
        )
        # VisualizationRuntime publishes grpc_url only after this method
        # returns, so a late subscriber cannot race the initial StoreInfo and
        # blueprint through the SDK micro-batcher.
        self._recording.flush(timeout_sec=5.0)

    def process(self, envelope: DecodedDataPlaneView) -> None:
        self._apply_latest_settings()
        self._set_message_time(envelope)

        if isinstance(envelope, DecodedPointCloudView):
            self._publish_points(envelope)
            point_frame = True
        elif isinstance(envelope, DecodedPoseView):
            self._publish_pose(envelope)
            point_frame = False
        else:
            return

        published_ns = time.monotonic_ns()
        decode_publish_ms = (
            published_ns - envelope.context.processing_started_monotonic_ns
        ) / 1_000_000
        if point_frame:
            self.metrics.published_pcl(
                self._last_point_count,
                published_ns,
                decode_publish_ms,
            )
        else:
            self.metrics.published_pose(
                published_ns,
                decode_publish_ms,
                len(self._path),
            )
        context = envelope.context
        if context.live and context.received_monotonic_ns is not None:
            self.metrics.record_latency((published_ns - context.received_monotonic_ns) / 1_000_000)

    def process_perception(self, frame: LivePerceptionResultFrame) -> None:
        """Publish one validated worker result on the live scene timeline."""

        self._apply_latest_settings()
        self._recording.set_time(
            "stream_time",
            timestamp=frame.captured_at_epoch_ns / 1_000_000_000,
        )
        self._recording.set_time(
            "capture_time",
            timestamp=frame.captured_at_epoch_ns / 1_000_000_000,
        )
        self._recording.set_time("message_sequence", sequence=frame.source_frame_index)
        self._recording.log(
            "/perception/camera/image",
            rr.EncodedImage(contents=frame.image_jpeg, media_type="image/jpeg"),
        )
        if frame.segmentation_mask is None:
            self._recording.log(
                "/perception/camera/segmentation",
                rr.Clear(recursive=False),
            )
        else:
            self._recording.log(
                "/perception/camera/segmentation",
                rr.SegmentationImage(frame.segmentation_mask),
            )
        if frame.objects:
            self._recording.log(
                "/perception/camera/detections",
                rr.Boxes2D(
                    array=[item["bbox_xyxy"] for item in frame.objects],
                    array_format=rr.Box2DFormat.XYXY,
                    labels=[_perception_label(item) for item in frame.objects],
                    colors=[
                        _perception_color(str(item["label"]), alpha=255)
                        for item in frame.objects
                    ],
                    show_labels=True,
                ),
            )
        else:
            self._recording.log(
                "/perception/camera/detections",
                rr.Clear(recursive=False),
            )
        cuboids = [
            item for item in frame.objects if item.get("cuboid_center_map") is not None
        ]
        if cuboids:
            self._recording.log(
                "/world/perception/boxes3d",
                rr.Boxes3D(
                    centers=[item["cuboid_center_map"] for item in cuboids],
                    half_sizes=[item["cuboid_half_size"] for item in cuboids],
                    quaternions=[
                        rr.Quaternion(xyzw=item["cuboid_quaternion_xyzw"])
                        for item in cuboids
                    ],
                    colors=[_perception_color(str(item["label"]), alpha=96) for item in cuboids],
                    labels=[_perception_label(item) for item in cuboids],
                    fill_mode=FillMode.Solid,
                    show_labels=True,
                ),
            )
        else:
            self._recording.log(
                "/world/perception/boxes3d",
                rr.Clear(recursive=False),
            )
        self.metrics.published_perception(
            captured_at_epoch_ns=frame.captured_at_epoch_ns,
            published_at_epoch_ns=time.time_ns(),
            published_monotonic_ns=time.monotonic_ns(),
        )

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        recording = self._recording
        try:
            try:
                recording.flush(timeout_sec=5.0)
            finally:
                recording.disconnect()
        finally:
            # The Python wrapper owns the native gRPC server. Release it even if
            # flush or disconnect fails instead of waiting for garbage collection.
            del self._recording

    def _set_message_time(self, envelope: DecodedDataPlaneView) -> None:
        context = envelope.context
        self._recording.set_time("stream_time", timestamp=time.time())
        self._recording.set_time(
            "capture_time",
            timestamp=context.captured_at_epoch_ns / 1_000_000_000,
        )
        self._recording.set_time("message_sequence", sequence=context.sequence)

    def _apply_latest_settings(self) -> None:
        settings = self._settings_provider()
        if settings == self._settings:
            return
        self._settings = settings
        self._recording.send_blueprint(_blueprint(settings))

    def _publish_points(self, frame: DecodedPointCloudView) -> None:
        positions = np.asarray(frame.positions_xyz, dtype=np.float32).reshape((-1, 3))
        if frame.intensities is None:
            intensities = np.full(frame.point_count, 255, dtype=np.uint8)
        else:
            intensities = np.frombuffer(frame.intensities, dtype=np.uint8)
        rgb = (
            None
            if frame.colors_rgb is None
            else np.frombuffer(frame.colors_rgb, dtype=np.uint8).reshape((-1, 3))
        )
        self._last_point_count = int(positions.shape[0])
        if not self._settings.show_points:
            self._recording.log("/world/points", rr.Clear(recursive=False))
            return
        colors = _point_colors(positions, intensities, rgb, self._settings)
        self._recording.log(
            "/world/points",
            rr.Points3D(
                positions,
                colors=colors,
                radii=rr.Radius.ui_points(self._settings.point_size),
            ),
        )

    def _publish_pose(self, frame: DecodedPoseView) -> None:
        publish_now_ns = time.monotonic_ns()
        position = (
            float(frame.position_xyz[0]),
            float(frame.position_xyz[1]),
            float(frame.position_xyz[2]),
        )
        self._recording.log(
            "/world/sensor_pose",
            rr.Transform3D(
                translation=position,
                quaternion=rr.Quaternion(xyzw=frame.orientation_xyzw),
            ),
        )

        appended = self._append_trajectory_pose(
            position,
            frame.context.captured_at_epoch_ns,
        )
        if not self._settings.show_trajectory:
            self._recording.log("/world/trajectory", rr.Clear(recursive=False))
            return
        if not appended:
            return
        if (
            publish_now_ns - self._last_trajectory_publish_ns
            < TRAJECTORY_PUBLISH_INTERVAL_NS
        ):
            return
        self._last_trajectory_publish_ns = publish_now_ns
        self._recording.log(
            "/world/trajectory",
            rr.LineStrips3D(
                [list(self._path)],
                colors=[247, 248, 244, 255],
                radii=rr.Radius.ui_points(2.0),
            ),
        )

    def _append_trajectory_pose(
        self,
        position: tuple[float, float, float],
        source_time_ns: int,
    ) -> bool:
        if not self._path:
            self._path.append(position)
            self._last_path_append_source_ns = source_time_ns
            return True

        elapsed_ns = source_time_ns - self._last_path_append_source_ns
        if elapsed_ns < 0:
            # A source clock discontinuity must not freeze trajectory sampling
            # until the timestamp catches up again. Preserve message order and
            # start a new sampling interval at the discontinuity.
            elapsed_ns = TRAJECTORY_FORCE_APPEND_NS
        if elapsed_ns < TRAJECTORY_APPEND_INTERVAL_NS:
            return False
        if (
            math.dist(self._path[-1], position) < TRAJECTORY_MIN_DISTANCE_METERS
            and elapsed_ns < TRAJECTORY_FORCE_APPEND_NS
        ):
            return False

        self._path.append(position)
        self._last_path_append_source_ns = source_time_ns
        if len(self._path) > MAX_TRAJECTORY_POSES:
            last = self._path[-1]
            self._path = self._path[::2]
            if self._path[-1] != last:
                self._path.append(last)
        return True


def _blueprint(settings: RerunSceneSettings) -> rrb.Blueprint:
    accumulation = max(0.0, settings.accumulation_seconds)
    time_range = rr.VisibleTimeRange(
        "stream_time",
        start=rr.TimeRangeBoundary.cursor_relative(seconds=-accumulation),
        end=rr.TimeRangeBoundary.cursor_relative(),
    )
    perception_active = (
        settings.show_detections_2d
        or settings.show_segmentation
        or settings.show_cuboids_3d
    )
    spatial = rrb.Spatial3DView(
        origin="/world",
        name="Мир · LiDAR и объекты" if perception_active else "Пространственная сцена",
        background=[7, 8, 10, 255],
        line_grid=rrb.LineGrid3D(
            visible=settings.show_grid,
            color=[86, 91, 99, 110],
            stroke_width=0.75,
        ),
        time_ranges=[] if perception_active else [time_range],
    )
    spatial.visualizer_overrides["/world/points"] = rrb.EntityBehavior(
        visible=settings.show_points
    )
    spatial.visualizer_overrides["/world/trajectory"] = rrb.EntityBehavior(
        visible=settings.show_trajectory
    )
    spatial.visualizer_overrides["/world/perception/boxes3d"] = rrb.EntityBehavior(
        visible=settings.show_cuboids_3d
    )
    camera = rrb.Spatial2DView(
        origin="/perception/camera",
        name="Оригинальное видео · слои AI",
    )
    camera.visualizer_overrides["/perception/camera/image"] = rrb.EntityBehavior(
        visible=True
    )
    camera.visualizer_overrides["/perception/camera/detections"] = rrb.EntityBehavior(
        visible=settings.show_detections_2d
    )
    camera.visualizer_overrides["/perception/camera/segmentation"] = rrb.EntityBehavior(
        visible=settings.show_segmentation
    )
    root = (
        rrb.Horizontal(camera, spatial, column_shares=[0.46, 0.54])
        if perception_active
        else spatial
    )
    return rrb.Blueprint(
        root,
        _live_time_panel(),
        auto_layout=False,
        auto_views=False,
        collapse_panels=True,
    )


def _perception_label(item: Mapping[str, Any]) -> str:
    base = f"#{int(item['track_id'])} {item['label']} · {float(item['score']):.0%}"
    distance = item.get("distance_m")
    return base if distance is None else f"{base} · {float(distance):.1f} m"


def _perception_color(label: str, *, alpha: int) -> list[int]:
    colors = {
        "person": (255, 99, 132),
        "car": (64, 180, 255),
        "truck": (255, 180, 64),
        "bus": (255, 210, 64),
        "bicycle": (110, 240, 155),
        "motorcycle": (170, 115, 255),
    }
    red, green, blue = colors.get(label, (247, 248, 244))
    return [red, green, blue, alpha]


def _live_time_panel() -> rrb.TimePanel:
    """Keep the hidden vendor timeline on its native live edge."""
    return rrb.TimePanel(
        timeline="stream_time",
        play_state="following",
        state="hidden",
    )


def _point_colors(
    positions: np.ndarray,
    intensities: np.ndarray,
    rgb: np.ndarray | None,
    settings: RerunSceneSettings,
) -> np.ndarray:
    if settings.palette == "custom":
        color = _parse_hex_color(settings.custom_color)
        return np.tile(np.asarray(color, dtype=np.uint8), (positions.shape[0], 1))
    if settings.color_mode == "rgb" and rgb is not None:
        return rgb
    if settings.color_mode == "class":
        color = _parse_hex_color(settings.custom_color)
        return np.tile(np.asarray(color, dtype=np.uint8), (positions.shape[0], 1))
    if settings.color_mode == "height":
        values = _normalize(positions[:, 2])
    elif settings.color_mode == "distance":
        values = _normalize(np.linalg.norm(positions, axis=1))
    else:
        values = intensities.astype(np.float32) / 255.0
    return _apply_palette(values, settings.palette, settings.custom_color)


def _normalize(values: np.ndarray) -> np.ndarray:
    if values.size == 0:
        return values.astype(np.float32)
    minimum = float(np.min(values))
    maximum = float(np.max(values))
    if not math.isfinite(minimum) or not math.isfinite(maximum) or maximum <= minimum:
        return np.full(values.shape, 0.5, dtype=np.float32)
    return ((values - minimum) / (maximum - minimum)).astype(np.float32)


def _apply_palette(values: np.ndarray, palette: PointPalette, custom: str) -> np.ndarray:
    clipped = np.clip(values, 0.0, 1.0)
    if palette == "custom":
        color = _parse_hex_color(custom)
        return np.tile(np.asarray(color, dtype=np.uint8), (values.size, 1))
    if palette == "grayscale":
        channel = np.rint(clipped * 255).astype(np.uint8)
        return np.column_stack((channel, channel, channel))

    stops = {
        "turbo": (
            (0.00, (48, 18, 59)),
            (0.25, (40, 126, 231)),
            (0.50, (42, 240, 154)),
            (0.75, (246, 214, 55)),
            (1.00, (180, 4, 38)),
        ),
        "viridis": (
            (0.00, (68, 1, 84)),
            (0.25, (59, 82, 139)),
            (0.50, (33, 145, 140)),
            (0.75, (94, 201, 98)),
            (1.00, (253, 231, 37)),
        ),
        "plasma": (
            (0.00, (13, 8, 135)),
            (0.25, (126, 3, 168)),
            (0.50, (204, 71, 120)),
            (0.75, (248, 149, 64)),
            (1.00, (240, 249, 33)),
        ),
    }[palette]
    positions = np.asarray([item[0] for item in stops], dtype=np.float32)
    colors = np.asarray([item[1] for item in stops], dtype=np.float32)
    channels = [np.interp(clipped, positions, colors[:, index]) for index in range(3)]
    return np.rint(np.column_stack(channels)).astype(np.uint8)


def _parse_hex_color(value: str) -> tuple[int, int, int]:
    candidate = value.removeprefix("#")
    if len(candidate) != 6:
        return 247, 248, 244
    try:
        return tuple(int(candidate[index : index + 2], 16) for index in (0, 2, 4))  # type: ignore[return-value]
    except ValueError:
        return 247, 248, 244