from __future__ import annotations

import json
import struct
import threading
import time
from pathlib import Path

import numpy as np
import pytest

from k1link.compute.live_perception import LivePerceptionResultFrame
from k1link.data_plane import DecodedDataPlaneView, NormalizationError
from k1link.device_plugins.xgrids_k1.mqtt.capture import FRAME_HEADER, RAW_MAGIC
from k1link.device_plugins.xgrids_k1.protocol.normalizer import normalize_k1_message
from k1link.device_plugins.xgrids_k1.viewer.messages import StreamMessage
from k1link.device_plugins.xgrids_k1.viewer.runtime import VisualizationRuntime
from k1link.viewer.rerun_bridge import (
    RerunBridge,
    RerunSceneSettings,
    _live_time_panel,
    _point_colors,
)


class FakeRecording:
    def __init__(self) -> None:
        self.logs: list[tuple[str, object, bool]] = []
        self.times: list[tuple[str, dict[str, object]]] = []
        self.blueprints: list[object] = []
        self.disconnected = False
        self.flush_count = 0

    def serve_grpc(self, **_: object) -> str:
        return "rerun+http://127.0.0.1:9876/proxy"

    def log(self, path: str, entity: object, *, static: bool = False) -> None:
        self.logs.append((path, entity, static))

    def set_time(self, timeline: str, **value: object) -> None:
        self.times.append((timeline, value))

    def send_blueprint(self, blueprint: object, **_: object) -> None:
        self.blueprints.append(blueprint)

    def disconnect(self) -> None:
        self.disconnected = True

    def flush(self, **_: object) -> None:
        self.flush_count += 1


class BlueprintFailureRecording(FakeRecording):
    def send_blueprint(self, blueprint: object, **kwargs: object) -> None:
        super().send_blueprint(blueprint, **kwargs)
        raise RuntimeError("synthetic blueprint failure")


class DisconnectFailureRecording(FakeRecording):
    def disconnect(self) -> None:
        super().disconnect()
        raise RuntimeError("synthetic disconnect failure")


def _message(
    topic: str,
    payload: bytes,
    *,
    sequence: int = 7,
    received_at_epoch_ns: int = 1_784_124_315_186_225_000,
) -> StreamMessage:
    return StreamMessage(
        sequence=sequence,
        topic=topic,
        payload=payload,
        received_at_epoch_ns=received_at_epoch_ns,
        received_monotonic_ns=None,
        source="test",
    )


def _envelope(
    topic: str,
    payload: bytes,
    *,
    sequence: int = 7,
    received_at_epoch_ns: int = 1_784_124_315_186_225_000,
) -> DecodedDataPlaneView:
    envelope = normalize_k1_message(
        _message(
            topic,
            payload,
            sequence=sequence,
            received_at_epoch_ns=received_at_epoch_ns,
        ),
        processing_started_monotonic_ns=time.monotonic_ns(),
    )
    assert envelope is not None
    return envelope


def test_legacy_points_and_pose_are_logged_to_rerun() -> None:
    recording = FakeRecording()
    bridge = RerunBridge(recording_factory=lambda _: recording)  # type: ignore[arg-type]

    point_payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    pose_payload = struct.pack("<ffffffff", 1.0, 2.0, 3.0, 99.0, 0.9, 0.1, 0.2, 0.3)
    bridge.process(_envelope("RealtimePointcloud", point_payload))
    bridge.process(_envelope("RealtimePath", pose_payload, sequence=8))

    paths = [path for path, _, _ in recording.logs]
    snapshot = bridge.metrics.snapshot()
    assert bridge.grpc_url == "rerun+http://127.0.0.1:9876/proxy"
    assert "/world/points" in paths
    assert "/world/sensor_pose" in paths
    assert "/world/trajectory" in paths
    assert snapshot["pcl_frames"] == 1
    assert snapshot["pose_frames"] == 1
    assert snapshot["last_point_count"] == 1
    assert {timeline for timeline, _ in recording.times} == {
        "capture_time",
        "message_sequence",
        "stream_time",
    }
    assert recording.flush_count == 1

    bridge.close()
    assert recording.disconnected is True
    assert recording.flush_count == 2


def test_live_blueprint_follows_stream_time_without_frontend_cursor_writes() -> None:
    panel = _live_time_panel()

    assert panel.timeline == "stream_time"
    assert panel.play_state == "following"
    assert panel.state == "hidden"


def test_live_perception_logs_original_mask_2d_distance_and_3d_cuboid() -> None:
    recording = FakeRecording()
    bridge = RerunBridge(recording_factory=lambda _: recording)  # type: ignore[arg-type]
    mask = np.zeros((600, 800), dtype=np.uint8)
    mask[20:30, 40:50] = 4

    bridge.process_perception(
        LivePerceptionResultFrame(
            frame_index=3,
            source_frame_index=30,
            session_seconds=1.0,
            captured_at_epoch_ns=1_784_124_315_186_225_000,
            image_jpeg=b"\xff\xd8test\xff\xd9",
            segmentation_mask=mask,
            objects=(
                {
                    "track_id": 9,
                    "label": "car",
                    "score": 0.9,
                    "bbox_xyxy": [10.0, 20.0, 110.0, 80.0],
                    "distance_m": 6.2,
                    "cuboid_center_map": [1.0, 2.0, 0.5],
                    "cuboid_half_size": [2.25, 0.925, 0.775],
                    "cuboid_quaternion_xyzw": [0.0, 0.0, 0.0, 1.0],
                },
            ),
            delivery={"health": "healthy"},
        )
    )

    paths = [path for path, _, _ in recording.logs]
    assert "/perception/camera/image" in paths
    assert "/perception/camera/segmentation" in paths
    assert "/perception/camera/detections" in paths
    assert "/world/perception/boxes3d" in paths


def test_constructor_disconnects_recording_after_partial_setup_failure() -> None:
    recording = BlueprintFailureRecording()

    with pytest.raises(RuntimeError, match="synthetic blueprint failure"):
        RerunBridge(recording_factory=lambda _: recording)  # type: ignore[arg-type]

    assert recording.disconnected is True


def test_fast_replay_trajectory_sampling_uses_source_time() -> None:
    recording = FakeRecording()
    bridge = RerunBridge(recording_factory=lambda _: recording)  # type: ignore[arg-type]
    base_time_ns = 1_784_124_315_000_000_000

    for index in range(4):
        pose_payload = struct.pack(
            "<ffffffff",
            index * 0.1,
            2.0,
            3.0,
            99.0,
            0.9,
            0.1,
            0.2,
            0.3,
        )
        bridge.process(
            _envelope(
                "RealtimePath",
                pose_payload,
                sequence=index + 1,
                received_at_epoch_ns=base_time_ns + index * 600_000_000,
            )
        )

    assert bridge.metrics.snapshot()["trajectory_poses"] == 4
    bridge.close()


def test_bad_frame_is_rejected_before_rerun_without_publishing() -> None:
    recording = FakeRecording()
    bridge = RerunBridge(recording_factory=lambda _: recording)  # type: ignore[arg-type]

    with pytest.raises(NormalizationError):
        normalize_k1_message(
            _message("RealtimePointcloud", b"short"),
            processing_started_monotonic_ns=time.monotonic_ns(),
        )

    snapshot = bridge.metrics.snapshot()
    assert snapshot["messages_received"] == 0
    assert snapshot["decode_errors"] == 0
    assert not any(path == "/world/points" for path, _, _ in recording.logs)


def test_palettes_are_deterministic_and_custom_color_is_exact() -> None:
    positions = np.asarray([[0, 0, 0], [0, 0, 10]], dtype=np.float32)
    intensities = np.asarray([0, 255], dtype=np.uint8)

    height = _point_colors(
        positions,
        intensities,
        None,
        RerunSceneSettings(color_mode="height", palette="viridis"),
    )
    custom = _point_colors(
        positions,
        intensities,
        None,
        RerunSceneSettings(color_mode="class", custom_color="#102030"),
    )
    custom_over_rgb = _point_colors(
        positions,
        intensities,
        np.asarray([[255, 0, 0], [0, 255, 0]], dtype=np.uint8),
        RerunSceneSettings(
            color_mode="rgb",
            palette="custom",
            custom_color="#102030",
        ),
    )

    assert height.shape == (2, 3)
    assert height.dtype == np.uint8
    assert height[0].tolist() == [68, 1, 84]
    assert height[1].tolist() == [253, 231, 37]
    assert custom.tolist() == [[16, 32, 48], [16, 32, 48]]
    assert custom_over_rgb.tolist() == [[16, 32, 48], [16, 32, 48]]


def test_runtime_reuses_one_bridge_across_sequential_sessions(tmp_path: Path) -> None:
    capture = tmp_path / "mqtt.raw.k1mqtt"
    point_topic = "RealtimePointcloud"
    pose_topic = "RealtimePath"
    point_payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    pose_payload = struct.pack("<ffffffff", 1.0, 2.0, 3.0, 99.0, 0.9, 0.1, 0.2, 0.3)
    frames = bytearray(RAW_MAGIC)
    for topic, payload in ((point_topic, point_payload), (pose_topic, pose_payload)):
        topic_raw = topic.encode()
        frames.extend(FRAME_HEADER.pack(len(topic_raw), len(payload)))
        frames.extend(topic_raw)
        frames.extend(payload)
    capture.write_bytes(frames)
    metadata = [
        {
            "record_type": "message",
            "sequence": 1,
            "received_at_epoch_ns": 1_000_000_000,
            "received_monotonic_ns": 1_000_000_000,
        },
        {
            "record_type": "message",
            "sequence": 2,
            "received_at_epoch_ns": 11_000_000_000,
            "received_monotonic_ns": 11_000_000_000,
        },
    ]
    (tmp_path / "mqtt.metadata.jsonl").write_text(
        "".join(json.dumps(item) + "\n" for item in metadata),
        encoding="utf-8",
    )

    recording = FakeRecording()
    created: list[RerunBridge] = []

    def bridge_factory(**kwargs: object) -> RerunBridge:
        bridge = RerunBridge(
            recording_factory=lambda _: recording,  # type: ignore[arg-type]
            **kwargs,  # type: ignore[arg-type]
        )
        created.append(bridge)
        return bridge

    runtime = VisualizationRuntime(
        bridge_factory=bridge_factory,
        normalizer=normalize_k1_message,
    )
    runtime.start_replay(capture, speed=1.0)
    deadline = time.monotonic() + 5.0
    snapshot = runtime.snapshot()
    while time.monotonic() < deadline:
        snapshot = runtime.snapshot()
        if snapshot["rerun_grpc_url"] and snapshot["metrics"]["pcl_frames"] == 1:
            break
        time.sleep(0.05)

    assert snapshot["rerun_grpc_url"] == "rerun+http://127.0.0.1:9876/proxy"
    assert snapshot["metrics"]["pcl_frames"] == 1
    assert snapshot["metrics"]["mqtt_to_publish_ms"] is None
    runtime.stop(wait_seconds=5.0)
    assert runtime.snapshot()["phase"] == "idle"
    assert runtime.snapshot()["rerun_grpc_url"] == "rerun+http://127.0.0.1:9876/proxy"
    assert recording.disconnected is False

    runtime.start_replay(capture, speed=0.0)
    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline:
        if runtime.snapshot()["metrics"]["pcl_frames"] == 1:
            break
        time.sleep(0.05)
    runtime.stop(wait_seconds=5.0)

    assert len(created) == 1
    assert runtime.snapshot()["metrics"]["pcl_frames"] == 1
    assert runtime.snapshot()["rerun_grpc_url"] == "rerun+http://127.0.0.1:9876/proxy"
    assert recording.disconnected is False
    runtime.close()
    assert runtime.snapshot()["rerun_grpc_url"] is None
    assert recording.disconnected is True


def test_runtime_reports_bridge_close_failure_instead_of_false_idle(tmp_path: Path) -> None:
    capture = tmp_path / "mqtt.raw.k1mqtt"
    topic = b"RealtimePointcloud"
    payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    capture.write_bytes(RAW_MAGIC + FRAME_HEADER.pack(len(topic), len(payload)) + topic + payload)
    (tmp_path / "mqtt.metadata.jsonl").write_text(
        json.dumps(
            {
                "record_type": "message",
                "sequence": 1,
                "received_at_epoch_ns": 1_000_000_000,
                "received_monotonic_ns": 1_000_000_000,
            }
        )
        + "\n",
        encoding="utf-8",
    )
    recording = DisconnectFailureRecording()

    runtime = VisualizationRuntime(
        bridge_factory=lambda **kwargs: RerunBridge(
            recording_factory=lambda _: recording,  # type: ignore[arg-type]
            **kwargs,  # type: ignore[arg-type]
        ),
        normalizer=normalize_k1_message,
    )
    runtime.start_replay(capture, speed=0.0)

    deadline = time.monotonic() + 5.0
    snapshot = runtime.snapshot()
    while snapshot["phase"] not in {"idle", "error"} and time.monotonic() < deadline:
        time.sleep(0.01)
        snapshot = runtime.snapshot()

    assert snapshot["phase"] == "idle"
    assert snapshot["rerun_grpc_url"] == "rerun+http://127.0.0.1:9876/proxy"
    assert recording.disconnected is False

    with pytest.raises(RuntimeError, match="synthetic disconnect failure"):
        runtime.close()

    snapshot = runtime.snapshot()
    assert snapshot["phase"] == "error"
    assert "synthetic disconnect failure" in snapshot["message"]
    assert snapshot["rerun_grpc_url"] is None
    assert recording.disconnected is True


def test_close_during_blocked_factory_closes_the_late_bridge(tmp_path: Path) -> None:
    capture = tmp_path / "mqtt.raw.k1mqtt"
    topic = b"RealtimePointcloud"
    payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    capture.write_bytes(RAW_MAGIC + FRAME_HEADER.pack(len(topic), len(payload)) + topic + payload)
    (tmp_path / "mqtt.metadata.jsonl").write_text(
        json.dumps(
            {
                "record_type": "message",
                "sequence": 1,
                "received_at_epoch_ns": 1_000_000_000,
                "received_monotonic_ns": 1_000_000_000,
            }
        )
        + "\n",
        encoding="utf-8",
    )

    factory_entered = threading.Event()
    release_factory = threading.Event()
    recording = FakeRecording()

    def blocked_factory(**kwargs: object) -> RerunBridge:
        factory_entered.set()
        assert release_factory.wait(timeout=5.0)
        return RerunBridge(
            recording_factory=lambda _: recording,  # type: ignore[arg-type]
            **kwargs,  # type: ignore[arg-type]
        )

    runtime = VisualizationRuntime(
        bridge_factory=blocked_factory,
        normalizer=normalize_k1_message,
    )
    runtime.start_replay(capture, speed=0.0)
    assert factory_entered.wait(timeout=2.0)

    with pytest.raises(RuntimeError, match="evidence lease сохранены"):
        runtime.close(wait_seconds=0.01)
    release_factory.set()
    deadline = time.monotonic() + 2.0
    while time.monotonic() < deadline and not recording.disconnected:
        time.sleep(0.01)
    runtime.close(wait_seconds=1.0)

    assert recording.disconnected is True
    assert runtime.snapshot()["rerun_grpc_url"] is None
    with pytest.raises(RuntimeError, match="runtime завершён"):
        runtime.start_replay(capture, speed=0.0)


def test_stop_fails_closed_when_runtime_thread_misses_deadline(tmp_path: Path) -> None:
    capture = tmp_path / "mqtt.raw.k1mqtt"
    topic = b"RealtimePointcloud"
    payload = struct.pack("<III", 16, 0, 0) + struct.pack(
        "<fffBBBB", 1.0, -2.0, 3.0, 10, 20, 30, 40
    )
    capture.write_bytes(RAW_MAGIC + FRAME_HEADER.pack(len(topic), len(payload)) + topic + payload)
    (tmp_path / "mqtt.metadata.jsonl").write_text(
        json.dumps(
            {
                "record_type": "message",
                "sequence": 1,
                "received_at_epoch_ns": 1_000_000_000,
                "received_monotonic_ns": 1_000_000_000,
            }
        )
        + "\n",
        encoding="utf-8",
    )

    factory_entered = threading.Event()
    release_factory = threading.Event()

    def blocked_factory(**kwargs: object) -> RerunBridge:
        factory_entered.set()
        assert release_factory.wait(timeout=5.0)
        return RerunBridge(
            recording_factory=lambda _: FakeRecording(),  # type: ignore[arg-type]
            **kwargs,  # type: ignore[arg-type]
        )

    runtime = VisualizationRuntime(
        bridge_factory=blocked_factory,
        normalizer=normalize_k1_message,
    )
    runtime.start_replay(capture, speed=0.0)
    assert factory_entered.wait(timeout=2.0)

    with pytest.raises(RuntimeError, match="не завершился"):
        runtime.stop(wait_seconds=0.01)
    assert runtime.snapshot()["phase"] == "stopping"

    release_factory.set()
    runtime.stop(wait_seconds=2.0)
    assert runtime.snapshot()["source_mode"] == "idle"
    runtime.close()