"""On-demand point-color overlays for prepared K1 recordings."""

from __future__ import annotations

import hashlib
import json
import os
import threading
from collections import OrderedDict
from contextlib import suppress
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

import numpy as np
import rerun as rr

from k1link.data_plane import DecodedPointCloudView, NormalizationError
from k1link.device_plugins.xgrids_k1.mqtt.capture import (
    CaptureFormatError,
    read_capture_clock_envelope,
    read_capture_clock_origin,
)
from k1link.device_plugins.xgrids_k1.protocol.normalizer import normalize_k1_message
from k1link.device_plugins.xgrids_k1.rrd_export import (
    APPLICATION_ID,
    CAPTURE_TIMELINE,
    SESSION_TIMELINE,
    _recorded_view_points,
    _should_publish_recorded_point_frame,
)
from k1link.device_plugins.xgrids_k1.viewer.messages import StreamMessage
from k1link.sessions import ReplayCommand
from k1link.viewer.rerun_bridge import RerunSceneSettings, _point_colors

_POINT_TOPICS = frozenset({"RealtimePointcloud", "lixel/application/report/lio_pcl"})
_MAX_METADATA_LINE_BYTES = 1024 * 1024
_MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
_DEFAULT_INDEX_CACHE_BYTES = 384 * 1024 * 1024
_DEFAULT_PAYLOAD_CACHE_BYTES = 192 * 1024 * 1024


class RecordedPointColorError(RuntimeError):
    """A recorded point-color overlay could not be prepared safely."""


@dataclass(frozen=True, slots=True)
class _PointColorFrame:
    sequence: int
    session_time_ns: int
    capture_time_ns: int
    positions: np.ndarray
    intensities: np.ndarray
    rgb: np.ndarray | None

    @property
    def byte_length(self) -> int:
        return (
            int(self.positions.nbytes)
            + int(self.intensities.nbytes)
            + (0 if self.rgb is None else int(self.rgb.nbytes))
        )


@dataclass(frozen=True, slots=True)
class _PointColorIndex:
    source_identity: tuple[object, ...]
    frames: tuple[_PointColorFrame, ...]
    byte_length: int


class RecordedPointColorOverlayStore:
    """Build small color-only RRD streams and retain bounded reusable inputs.

    The base operator RRD already owns point positions. These overlays log only
    ``Points3D:colors`` at the exact same session timestamps, so changing a
    palette never duplicates the point geometry, trajectory, video, or AI data.
    """

    def __init__(
        self,
        *,
        index_cache_bytes: int = _DEFAULT_INDEX_CACHE_BYTES,
        payload_cache_bytes: int = _DEFAULT_PAYLOAD_CACHE_BYTES,
    ) -> None:
        if index_cache_bytes < 0 or payload_cache_bytes < 0:
            raise ValueError("recorded color cache budgets must be non-negative")
        self._index_cache_bytes = index_cache_bytes
        self._payload_cache_bytes = payload_cache_bytes
        self._lock = threading.RLock()
        self._session_locks: dict[str, threading.Lock] = {}
        self._indexes: OrderedDict[tuple[object, ...], _PointColorIndex] = OrderedDict()
        self._index_bytes = 0
        self._payloads: OrderedDict[tuple[object, ...], bytes] = OrderedDict()
        self._payload_bytes = 0

    def render(
        self,
        command: ReplayCommand,
        *,
        application_id: str,
        recording_id: str,
        color_mode: Literal["intensity", "height", "distance", "rgb", "class"],
        palette: Literal["turbo", "viridis", "plasma", "grayscale", "custom"],
        custom_color: str,
    ) -> bytes:
        if command.plugin_id != "nodedc.device.xgrids-lixelkity-k1":
            raise RecordedPointColorError("recorded color provider received another plugin")
        if application_id != APPLICATION_ID:
            raise ValueError("recorded color application id is invalid")
        if not recording_id or len(recording_id) > 128:
            raise ValueError("recorded color recording id is invalid")

        source = command.primary_artifact.path.expanduser().resolve()
        metadata = _artifact_path(command, "raw-transport-index")
        capture_clock = _artifact_path(command, "raw-transport-clock")
        capture_clock_origin = _artifact_path(command, "raw-transport-clock-origin")
        source_identity = _source_identity(
            source,
            metadata,
            capture_clock,
            capture_clock_origin,
        )
        settings = RerunSceneSettings(
            color_mode=color_mode,
            palette=palette,
            custom_color=custom_color,
        )
        settings_key = (
            color_mode,
            palette,
            (
                custom_color.casefold()
                if palette == "custom" or color_mode == "class"
                else "-"
            ),
        )
        payload_key = (*source_identity, application_id, recording_id, *settings_key)

        with self._lock:
            cached_payload = self._payloads.get(payload_key)
            if cached_payload is not None:
                self._payloads.move_to_end(payload_key)
                return cached_payload
            session_lock = self._session_locks.setdefault(command.session_id, threading.Lock())

        with session_lock:
            with self._lock:
                cached_payload = self._payloads.get(payload_key)
                if cached_payload is not None:
                    self._payloads.move_to_end(payload_key)
                    return cached_payload
            index = self._index(
                source_identity,
                source,
                metadata,
                capture_clock,
                capture_clock_origin,
            )
            payload = _render_color_overlay(
                index.frames,
                application_id=application_id,
                recording_id=recording_id,
                settings=settings,
            )
            self._remember_payload(payload_key, payload)
            return payload

    def _index(
        self,
        source_identity: tuple[object, ...],
        source: Path,
        metadata: Path | None,
        capture_clock: Path | None,
        capture_clock_origin: Path | None,
    ) -> _PointColorIndex:
        with self._lock:
            cached = self._indexes.get(source_identity)
            if cached is not None:
                self._indexes.move_to_end(source_identity)
                return cached
        if metadata is None:
            raise RecordedPointColorError("native point-color index is unavailable")
        index = _build_index(
            source,
            metadata,
            capture_clock,
            capture_clock_origin,
            source_identity=source_identity,
        )
        if index.byte_length > self._index_cache_bytes:
            return index
        with self._lock:
            existing = self._indexes.pop(source_identity, None)
            if existing is not None:
                self._index_bytes -= existing.byte_length
            self._indexes[source_identity] = index
            self._index_bytes += index.byte_length
            while self._indexes and self._index_bytes > self._index_cache_bytes:
                _, stale = self._indexes.popitem(last=False)
                self._index_bytes -= stale.byte_length
        return index

    def _remember_payload(self, key: tuple[object, ...], payload: bytes) -> None:
        if len(payload) > self._payload_cache_bytes:
            return
        with self._lock:
            existing = self._payloads.pop(key, None)
            if existing is not None:
                self._payload_bytes -= len(existing)
            self._payloads[key] = payload
            self._payload_bytes += len(payload)
            while self._payloads and self._payload_bytes > self._payload_cache_bytes:
                _, stale = self._payloads.popitem(last=False)
                self._payload_bytes -= len(stale)


def _build_index(
    source: Path,
    metadata: Path,
    capture_clock: Path | None,
    capture_clock_origin: Path | None,
    *,
    source_identity: tuple[object, ...],
) -> _PointColorIndex:
    try:
        envelope = None if capture_clock is None else read_capture_clock_envelope(capture_clock)
        origin = (
            None
            if capture_clock_origin is None
            else read_capture_clock_origin(capture_clock_origin)
        )
    except CaptureFormatError as exc:
        raise RecordedPointColorError("native point-color clock is invalid") from exc
    if envelope is not None and origin is not None and (
        envelope.started_at_epoch_ns != origin.started_at_epoch_ns
        or envelope.started_monotonic_ns != origin.started_monotonic_ns
    ):
        raise RecordedPointColorError("native point-color clocks do not match")
    session_origin_ns = (
        envelope.started_monotonic_ns
        if envelope is not None
        else origin.started_monotonic_ns
        if origin is not None
        else None
    )

    frames: list[_PointColorFrame] = []
    total_bytes = 0
    point_frame_number = 0
    previous_sequence = 0
    previous_monotonic_ns: int | None = None
    try:
        source_size = source.stat().st_size
        with source.open("rb") as raw_stream, metadata.open("r", encoding="utf-8") as index_stream:
            for raw_line in index_stream:
                if len(raw_line.encode("utf-8")) > _MAX_METADATA_LINE_BYTES:
                    raise RecordedPointColorError("native point-color metadata line is too large")
                record = json.loads(raw_line)
                sequence = record.get("sequence")
                monotonic_ns = record.get("received_monotonic_ns")
                capture_time_ns = record.get("received_at_epoch_ns")
                if (
                    record.get("record_type") != "message"
                    or not isinstance(sequence, int)
                    or sequence <= previous_sequence
                    or not isinstance(monotonic_ns, int)
                    or monotonic_ns < 0
                    or not isinstance(capture_time_ns, int)
                    or capture_time_ns < 0
                ):
                    raise RecordedPointColorError("native point-color metadata is invalid")
                if previous_monotonic_ns is not None and monotonic_ns < previous_monotonic_ns:
                    raise RecordedPointColorError("native point-color timeline decreases")
                previous_sequence = sequence
                previous_monotonic_ns = monotonic_ns
                if session_origin_ns is None:
                    session_origin_ns = monotonic_ns
                topic = record.get("topic")
                if topic not in _POINT_TOPICS:
                    continue
                point_frame_number += 1
                if not _should_publish_recorded_point_frame(point_frame_number):
                    continue
                payload_offset = record.get("raw_payload_offset")
                payload_bytes = record.get("payload_bytes")
                payload_sha256 = record.get("payload_sha256")
                if (
                    not isinstance(payload_offset, int)
                    or payload_offset < 8
                    or not isinstance(payload_bytes, int)
                    or payload_bytes < 0
                    or payload_bytes > _MAX_PAYLOAD_BYTES
                    or payload_offset + payload_bytes > source_size
                    or not isinstance(payload_sha256, str)
                    or len(payload_sha256) != 64
                ):
                    raise RecordedPointColorError("native point-color offset is invalid")
                raw_stream.seek(payload_offset)
                payload = raw_stream.read(payload_bytes)
                if (
                    len(payload) != payload_bytes
                    or hashlib.sha256(payload).hexdigest() != payload_sha256
                ):
                    raise RecordedPointColorError("native point-color payload is corrupt")
                message = StreamMessage(
                    sequence=sequence,
                    topic=topic,
                    payload=payload,
                    received_at_epoch_ns=capture_time_ns,
                    received_monotonic_ns=monotonic_ns,
                    source="k1mqtt",
                )
                try:
                    decoded = normalize_k1_message(
                        message,
                        processing_started_monotonic_ns=monotonic_ns,
                    )
                except NormalizationError as exc:
                    raise RecordedPointColorError(
                        "native point-color frame failed normalization"
                    ) from exc
                if not isinstance(decoded, DecodedPointCloudView):
                    raise RecordedPointColorError("native point-color topic is not a point frame")
                positions = np.asarray(decoded.positions_xyz, dtype=np.float32).reshape((-1, 3))
                intensities = (
                    np.full(decoded.point_count, 255, dtype=np.uint8)
                    if decoded.intensities is None
                    else np.frombuffer(decoded.intensities, dtype=np.uint8)
                )
                rgb = (
                    None
                    if decoded.colors_rgb is None
                    else np.frombuffer(decoded.colors_rgb, dtype=np.uint8).reshape((-1, 3))
                )
                positions, intensities, rgb = _recorded_view_points(
                    positions,
                    intensities,
                    rgb,
                )
                frame = _PointColorFrame(
                    sequence=sequence,
                    session_time_ns=monotonic_ns - session_origin_ns,
                    capture_time_ns=capture_time_ns,
                    positions=positions.copy(),
                    intensities=intensities.copy(),
                    rgb=None if rgb is None else rgb.copy(),
                )
                frames.append(frame)
                total_bytes += frame.byte_length
    except (OSError, json.JSONDecodeError, UnicodeError) as exc:
        raise RecordedPointColorError("native point-color index could not be read") from exc
    if session_origin_ns is None or not frames:
        raise RecordedPointColorError("native point-color index contains no point frames")
    return _PointColorIndex(
        source_identity=source_identity,
        frames=tuple(frames),
        byte_length=total_bytes,
    )


def _render_color_overlay(
    frames: tuple[_PointColorFrame, ...],
    *,
    application_id: str,
    recording_id: str,
    settings: RerunSceneSettings,
) -> bytes:
    recording = rr.RecordingStream(
        application_id,
        recording_id=recording_id,
        send_properties=False,
    )
    stream = rr.binary_stream(recording)
    try:
        for frame in frames:
            recording.set_time(
                SESSION_TIMELINE,
                duration=np.timedelta64(frame.session_time_ns, "ns"),
            )
            recording.set_time(
                CAPTURE_TIMELINE,
                timestamp=np.datetime64(frame.capture_time_ns, "ns"),
            )
            recording.set_time("message_sequence", sequence=frame.sequence)
            recording.log(
                "/world/points",
                rr.Points3D.from_fields(
                    colors=_point_colors(
                        frame.positions,
                        frame.intensities,
                        frame.rgb,
                        settings,
                    )
                ),
            )
        payload = stream.read(flush=True, flush_timeout_sec=30.0)
    except Exception as exc:
        raise RecordedPointColorError("recorded point-color RRD serialization failed") from exc
    finally:
        with suppress(Exception):
            recording.disconnect()
    if not payload or not payload.startswith(b"RRF2"):
        raise RecordedPointColorError("recorded point-color RRD is invalid")
    return payload


def _artifact_path(command: ReplayCommand, artifact_id: str) -> Path | None:
    matches = tuple(
        artifact.path.expanduser().resolve()
        for artifact in command.artifacts
        if artifact.artifact_id == artifact_id
    )
    if not matches:
        return None
    if len(matches) != 1:
        raise RecordedPointColorError("recorded point-color artifact is ambiguous")
    return matches[0]


def _source_identity(source: Path, *companions: Path | None) -> tuple[object, ...]:
    values: list[object] = []
    for path in (source, *companions):
        if path is None:
            values.extend((None, None, None))
            continue
        try:
            metadata = path.stat()
        except OSError as exc:
            raise RecordedPointColorError("recorded point-color source is unavailable") from exc
        if not os.path.isfile(path):
            raise RecordedPointColorError("recorded point-color source is not a file")
        values.extend((str(path), metadata.st_size, metadata.st_mtime_ns))
    return tuple(values)
