"""Vendor-neutral Rerun blueprint for prepared observation recordings."""

from __future__ import annotations

from collections import OrderedDict
from collections.abc import Callable
from contextlib import suppress
from threading import Lock
from typing import Any, Literal
from uuid import UUID

import rerun as rr
import rerun_bindings as bindings
from rerun import blueprint as rrb

from k1link.viewer.rerun_bridge import RerunSceneSettings

APPLICATION_ID = "nodedc_mission_core_recorded"
SESSION_TIMELINE = "session_time"

RECORDED_SPATIAL_VIEW_ID = UUID("5f5f11d5-3b0a-4a81-887b-2be767cba1c0")
RECORDED_SPATIAL_RESET_VIEW_ID = UUID("1480abcc-a0ae-4287-ab69-4b606d15d947")
RECORDED_ROOT_CONTAINER_ID = UUID("b02f2aca-8471-4dcb-b786-53df5a320fc8")
RECORDED_SPATIAL_RESET_ROOT_CONTAINER_ID = UUID("27d95ad7-1e72-46ca-b854-d02cd99b7610")
RECORDED_PERCEPTION_ROOT_CONTAINER_ID = UUID("70ea9fd5-bbbb-4b23-8ae8-8098af92b997")
RECORDED_PERCEPTION_RESET_ROOT_CONTAINER_ID = UUID("d2196c6e-99da-4402-857c-4daea0dd3ae0")
RECORDED_PERCEPTION_3D_ROOT_CONTAINER_ID = UUID("81051015-9808-413b-90a4-1cfaacacbc4f")
RECORDED_PERCEPTION_3D_RESET_ROOT_CONTAINER_ID = UUID("3de10822-0274-4a58-8d45-f35d86625303")
RECORDED_METRICS_ROOT_CONTAINER_ID = UUID("374710a2-1d77-4349-bea4-8719e7aa1f09")
RECORDED_METRICS_RESET_ROOT_CONTAINER_ID = UUID("1e8ac565-bbd6-4554-9213-dad564e534e3")
RECORDED_POINTS_VISUALIZER_ID = UUID("ca037ec0-8761-4417-86ee-846fa2875303")
RECORDED_PERCEPTION_POINTS_VISUALIZER_ID = UUID("ab8db306-6981-49c5-b417-94fe2f01dbf0")
RECORDED_CAMERA_VIEW_ID = UUID("5c1db75b-07cd-479a-903d-f4f2ed554513")
RECORDED_CAMERA_RESET_VIEW_ID = UUID("d0618e0c-c889-4222-a643-60a4a882935c")
RECORDED_PERCEPTION_3D_VIEW_ID = UUID("0496bd2e-2b4d-4a4f-87b8-3ce4f9f7e114")
RECORDED_PERCEPTION_3D_RESET_VIEW_ID = UUID("31f63ab8-ffdd-4751-ae40-e8e7b4462c8f")
RECORDED_METRICS_VIEW_ID = UUID("f973fc11-0867-4732-ad3c-97008621fab7")
RECORDED_UNIFIED_ROOT_CONTAINER_ID = UUID("e9934ef8-453f-432e-9136-2b0908190253")
RECORDED_UNIFIED_RESET_ROOT_CONTAINER_ID = UUID("2781407d-9f4e-4405-86e8-bbe6065e83df")
RecordedView = Literal["spatial", "perception", "perception3d", "metrics"]


class RecordedBlueprintError(RuntimeError):
    """A viewer blueprint update could not be serialized safely."""


class _RecordedBlueprintStream:
    """Keep one browser viewport on one mutable Rerun blueprint store.

    Rerun persists the operator-controlled eye inside the active blueprint
    store. Creating and activating a fresh store for every layer toggle drops
    that eye and snaps the view back to its fallback camera. This stream keeps
    the store identity stable until the operator explicitly requests a reset.
    """

    def __init__(
        self,
        application_id: str,
        recording_id: str,
        *,
        view_reset_generation: Literal[0, 1],
    ) -> None:
        self.view_reset_generation = view_reset_generation
        self._lock = Lock()
        self._sequence = 0
        self._follow_trajectory: bool | None = None
        self._closed = False
        native = bindings.new_blueprint(
            application_id=application_id,
            make_default=False,
            make_thread_default=False,
            default_enabled=True,
        )
        self._blueprint_recording = rr.RecordingStream._from_native(native)
        self._blueprint_memory = self._blueprint_recording.memory_recording()
        self._transport_recording = rr.RecordingStream(
            application_id,
            recording_id=recording_id,
            send_properties=False,
        )
        self._transport = rr.binary_stream(self._transport_recording)

    def render(
        self,
        blueprint_factory: Callable[[bool], rrb.Blueprint],
        *,
        follow_trajectory: bool,
    ) -> bytes:
        with self._lock:
            if self._closed:
                raise RecordedBlueprintError("stable blueprint stream is closed")
            update_eye_controls = (
                self._follow_trajectory is None
                or self._follow_trajectory != follow_trajectory
            )
            blueprint = blueprint_factory(update_eye_controls)
            self._blueprint_recording.set_time(
                "blueprint",
                sequence=self._sequence,
            )
            self._sequence += 1
            blueprint._log_to_stream(self._blueprint_recording)
            self._blueprint_recording.flush(timeout_sec=5.0)
            bindings.send_blueprint(
                self._blueprint_memory.storage,
                True,
                False,
                self._transport_recording.to_native(),
            )
            payload = self._transport.read(flush=True, flush_timeout_sec=5.0)
            if not payload:
                raise RecordedBlueprintError("stable blueprint stream produced no data")
            self._follow_trajectory = follow_trajectory
            return payload

    def close(self) -> None:
        with self._lock:
            self._closed = True
            with suppress(Exception):
                self._blueprint_recording.disconnect()
            with suppress(Exception):
                self._transport_recording.disconnect()


_MAX_RECORDED_BLUEPRINT_STREAMS = 32
_recorded_blueprint_streams_lock = Lock()
_recorded_blueprint_streams: OrderedDict[
    tuple[str, str, str], _RecordedBlueprintStream
] = OrderedDict()


def _stable_recorded_blueprint_stream(
    *,
    application_id: str,
    recording_id: str,
    blueprint_session_id: str,
    view_reset_generation: Literal[0, 1],
) -> _RecordedBlueprintStream:
    key = (application_id, recording_id, blueprint_session_id)
    with _recorded_blueprint_streams_lock:
        stream = _recorded_blueprint_streams.get(key)
        if stream is not None and stream.view_reset_generation != view_reset_generation:
            stream.close()
            del _recorded_blueprint_streams[key]
            stream = None
        if stream is None:
            stream = _RecordedBlueprintStream(
                application_id,
                recording_id,
                view_reset_generation=view_reset_generation,
            )
            _recorded_blueprint_streams[key] = stream
        else:
            _recorded_blueprint_streams.move_to_end(key)
        while len(_recorded_blueprint_streams) > _MAX_RECORDED_BLUEPRINT_STREAMS:
            _, stale = _recorded_blueprint_streams.popitem(last=False)
            stale.close()
        return stream


def recorded_blueprint(
    settings: RerunSceneSettings,
    *,
    include_initial_playback_state: bool = True,
    active_view: RecordedView = "spatial",
    view_reset_generation: Literal[0, 1] = 0,
    unified_perception: bool = False,
    show_detections_2d: bool = False,
    show_segmentation: bool = False,
    show_cuboids_3d: bool = False,
    follow_trajectory: bool = False,
    update_eye_controls: bool = True,
) -> rrb.Blueprint:
    accumulation = max(0.0, settings.accumulation_seconds)
    time_ranges: list[rr.VisibleTimeRange] | None = None
    if accumulation > 0:
        time_ranges = [
            rr.VisibleTimeRange(
                SESSION_TIMELINE,
                start=rr.TimeRangeBoundary.cursor_relative(seconds=-accumulation),
                end=rr.TimeRangeBoundary.cursor_relative(),
            )
        ]
    accumulated_time_ranges = (
        rrb.VisibleTimeRanges(time_ranges) if time_ranges is not None else None
    )
    point_visualizer = rr.Points3D.from_fields(
        radii=rr.Radius.ui_points(settings.point_size),
    ).visualizer()
    point_visualizer.id = RECORDED_POINTS_VISUALIZER_ID
    perception_point_visualizer = rr.Points3D.from_fields(
        radii=rr.Radius.ui_points(settings.point_size),
    ).visualizer()
    perception_point_visualizer.id = RECORDED_PERCEPTION_POINTS_VISUALIZER_ID
    point_overrides: list[Any] = [
        rrb.EntityBehavior(visible=settings.show_points),
        point_visualizer,
    ]
    trajectory_overrides: list[Any] = [
        rrb.EntityBehavior(visible=settings.show_trajectory),
    ]
    if accumulated_time_ranges is not None:
        point_overrides.append(accumulated_time_ranges)
        trajectory_overrides.append(accumulated_time_ranges)
    spatial_view = rrb.Spatial3DView(
        origin="/world",
        name="Мир · LiDAR и объекты" if unified_perception else "Пространственная сцена",
        background=[7, 8, 10, 255],
        line_grid=rrb.LineGrid3D(
            visible=settings.show_grid,
            color=[86, 91, 99, 110],
            stroke_width=0.75,
        ),
        overrides={
            # Accumulation belongs to mapping data only. Dynamic perception
            # inherits the view's latest-at query and can never turn into an
            # object-history trail when the operator widens the cloud window.
            "/world/points": point_overrides,
            "/world/trajectory": trajectory_overrides,
            "/world/perception": rrb.EntityBehavior(visible=show_cuboids_3d),
            "/world/perception/lidar": rrb.EntityBehavior(visible=False),
            "/world/perception/support": rrb.EntityBehavior(visible=False),
            "/world/perception/semantic_points": rrb.EntityBehavior(visible=False),
            "/world/perception/boxes3d": rrb.EntityBehavior(
                visible=show_cuboids_3d,
            ),
        },
        # Log an explicit empty range set so a previous accumulated blueprint
        # for this stable view id is cleared instead of surviving in Rerun.
        time_ranges=rrb.VisibleTimeRanges([]),
        # Follow is a property of the same operator eye, not another view.
        # Keeping the view id stable preserves the current orbit offset when
        # tracking is toggled. An explicit empty path clears tracking without
        # overwriting the position/look-target saved by user interaction.
        eye_controls=(
            rrb.EyeControls3D.from_fields(
                kind=rrb.Eye3DKind.Orbital if follow_trajectory else None,
                tracking_entity="/world/sensor_pose" if follow_trajectory else "",
            )
            if update_eye_controls
            else None
        ),
    )
    spatial_view.id = (
        RECORDED_SPATIAL_RESET_VIEW_ID
        if view_reset_generation
        else RECORDED_SPATIAL_VIEW_ID
    )
    camera_view = rrb.Spatial2DView(
        origin="/perception/camera",
        name="Оригинальное видео · слои AI",
        background=[7, 8, 10, 255],
        overrides={
            "/perception/camera/image": rrb.EntityBehavior(visible=True),
            "/perception/camera/detections": rrb.EntityBehavior(
                visible=show_detections_2d,
            ),
            "/perception/camera/segmentation": rrb.EntityBehavior(
                visible=show_segmentation,
            ),
        },
    )
    camera_view.id = (
        RECORDED_CAMERA_RESET_VIEW_ID if view_reset_generation else RECORDED_CAMERA_VIEW_ID
    )
    perception_3d_view = rrb.Spatial3DView(
        # Cuboids are expressed in the same calibrated world frame as the
        # native LiDAR recording.  Rooting this view at /world/perception
        # excluded /world/points and left the operator looking at boxes in an
        # empty scene.  Keep the main point cloud and the derived objects in
        # one latest-at world view; unlike the mapping tab, this deliberately
        # has no accumulated time range, so dynamic cuboids never stack.
        origin="/world",
        name="Сегментация и объекты · 3D",
        background=[7, 8, 10, 255],
        line_grid=rrb.LineGrid3D(
            visible=settings.show_grid,
            color=[86, 91, 99, 110],
            stroke_width=0.75,
        ),
        overrides={
            "/world/points": [
                rrb.EntityBehavior(visible=settings.show_points),
                perception_point_visualizer,
            ],
            "/world/trajectory": rrb.EntityBehavior(
                visible=settings.show_trajectory,
            ),
            "/world/perception": rrb.EntityBehavior(visible=True),
            # The overlay already carries the selected fusion support points;
            # its grey diagnostic LiDAR copy would otherwise double-render the
            # native cloud from /world/points.
            "/world/perception/lidar": rrb.EntityBehavior(visible=False),
        },
        eye_controls=(
            rrb.EyeControls3D.from_fields(
                kind=rrb.Eye3DKind.Orbital if follow_trajectory else None,
                tracking_entity="/world/sensor_pose" if follow_trajectory else "",
            )
            if update_eye_controls
            else None
        ),
    )
    perception_3d_view.id = (
        RECORDED_PERCEPTION_3D_RESET_VIEW_ID
        if view_reset_generation
        else RECORDED_PERCEPTION_3D_VIEW_ID
    )
    metrics_view = rrb.TimeSeriesView(
        origin="/metrics/device",
        name="Маршрут и время",
    )
    metrics_view.id = RECORDED_METRICS_VIEW_ID
    root_container: Any
    if unified_perception:
        # Operator perception is one synchronized composition, not a set of
        # mutually exclusive modes. The camera keeps the original frame as
        # its base and overlays 2D detections/segmentation. The paired world
        # view shows the same cursor in the native LiDAR frame and reveals
        # semantic points and cuboids independently.
        spatial_view.visualizer_overrides["/world/perception"] = [
            rrb.EntityBehavior(visible=True),
        ]
        spatial_view.visualizer_overrides["/world/perception/lidar"] = [
            rrb.EntityBehavior(visible=False),
        ]
        spatial_view.visualizer_overrides["/world/perception/support"] = [
            rrb.EntityBehavior(visible=False),
        ]
        spatial_view.visualizer_overrides["/world/perception/semantic_points"] = [
            rrb.EntityBehavior(visible=show_segmentation),
        ]
        spatial_view.visualizer_overrides["/world/perception/boxes3d"] = [
            rrb.EntityBehavior(visible=show_cuboids_3d),
        ]
        root_container = rrb.Horizontal(
            camera_view,
            spatial_view,
            column_shares=[0.46, 0.54],
            name="Единая сцена восприятия",
        )
        root_container.id = (
            RECORDED_UNIFIED_RESET_ROOT_CONTAINER_ID
            if view_reset_generation
            else RECORDED_UNIFIED_ROOT_CONTAINER_ID
        )
    else:
        # Keep operator video and 3D cuboids as direct root views. Rerun's nested
        # Tabs preserve their own active child and cannot be switched reliably by
        # a live blueprint channel after the operator has visited another child.
        active_tab = {"spatial": 0, "perception": 1, "perception3d": 2, "metrics": 3}[active_view]
        root_container = rrb.Tabs(
            spatial_view,
            camera_view,
            perception_3d_view,
            metrics_view,
            active_tab=active_tab,
        )
        # Rerun persists the active child on a Tabs container and does not reliably
        # replace it from a later blueprint message. Give every operator mode (and
        # its explicit reset generation) a stable root identity so the requested
        # child is authoritative instead of inheriting a previously visited tab.
        root_container.id = {
            ("spatial", 0): RECORDED_ROOT_CONTAINER_ID,
            ("spatial", 1): RECORDED_SPATIAL_RESET_ROOT_CONTAINER_ID,
            ("perception", 0): RECORDED_PERCEPTION_ROOT_CONTAINER_ID,
            ("perception", 1): RECORDED_PERCEPTION_RESET_ROOT_CONTAINER_ID,
            ("perception3d", 0): RECORDED_PERCEPTION_3D_ROOT_CONTAINER_ID,
            ("perception3d", 1): RECORDED_PERCEPTION_3D_RESET_ROOT_CONTAINER_ID,
            ("metrics", 0): RECORDED_METRICS_ROOT_CONTAINER_ID,
            ("metrics", 1): RECORDED_METRICS_RESET_ROOT_CONTAINER_ID,
        }[(active_view, view_reset_generation)]

    if include_initial_playback_state:
        return rrb.Blueprint(
            root_container,
            rrb.TimePanel(
                timeline=SESSION_TIMELINE,
                play_state="paused",
                state="hidden",
            ),
            auto_layout=False,
            auto_views=False,
            collapse_panels=True,
        )
    return rrb.Blueprint(
        root_container,
        auto_layout=False,
        auto_views=False,
        collapse_panels=False,
    )


def recorded_blueprint_rrd(
    settings: RerunSceneSettings,
    *,
    application_id: str = APPLICATION_ID,
    recording_id: str,
    blueprint_session_id: str | None = None,
    active_view: RecordedView = "spatial",
    view_reset_generation: Literal[0, 1] = 0,
    unified_perception: bool = False,
    show_detections_2d: bool = False,
    show_segmentation: bool = False,
    show_cuboids_3d: bool = False,
    follow_trajectory: bool = False,
) -> bytes:
    """Serialize a bounded active blueprint update without recorded data."""

    def build_blueprint(update_eye_controls: bool) -> rrb.Blueprint:
        return recorded_blueprint(
            settings,
            include_initial_playback_state=False,
            active_view=active_view,
            view_reset_generation=view_reset_generation,
            unified_perception=unified_perception,
            show_detections_2d=show_detections_2d,
            show_segmentation=show_segmentation,
            show_cuboids_3d=show_cuboids_3d,
            follow_trajectory=follow_trajectory,
            update_eye_controls=update_eye_controls,
        )

    payload: bytes | None
    if blueprint_session_id is not None:
        try:
            payload = _stable_recorded_blueprint_stream(
                application_id=application_id,
                recording_id=recording_id,
                blueprint_session_id=blueprint_session_id,
                view_reset_generation=view_reset_generation,
            ).render(
                build_blueprint,
                follow_trajectory=follow_trajectory,
            )
        except RecordedBlueprintError:
            raise
        except Exception as exc:
            raise RecordedBlueprintError("failed to serialize stable recorded blueprint") from exc
    else:
        blueprint = build_blueprint(True)
        recording = rr.RecordingStream(
            application_id,
            recording_id=recording_id,
            send_properties=False,
        )
        stream = rr.binary_stream(recording)
        try:
            recording.send_blueprint(
                blueprint,
                make_active=True,
                make_default=False,
            )
            payload = stream.read(flush=True, flush_timeout_sec=5.0)
        except Exception as exc:
            raise RecordedBlueprintError("failed to serialize recorded blueprint") from exc
        finally:
            with suppress(Exception):
                recording.disconnect()
    if not payload or not payload.startswith(b"RRF2") or len(payload) > 1_048_576:
        raise RecordedBlueprintError("serialized recorded blueprint is invalid")
    return payload
