from __future__ import annotations

import struct
import time
from collections import deque
from dataclasses import dataclass

import foxglove
from foxglove import Channel
from foxglove.channels import (
    PointCloudChannel,
    PoseInFrameChannel,
    SceneUpdateChannel,
)
from foxglove.messages import (
    Color,
    LinePrimitive,
    LinePrimitiveLineType,
    PackedElementField,
    PackedElementFieldNumericType,
    Point3,
    PointCloud,
    Pose,
    PoseInFrame,
    Quaternion,
    SceneEntity,
    SceneUpdate,
    Timestamp,
    Vector3,
)

from k1link.device_plugins.xgrids_k1.protocol.streams import (
    LegacyPointCloudFrame,
    LegacyPoseFrame,
    LioPointCloudFrame,
    LioPoseFrame,
    StreamDecodeError,
    decode_legacy_pointcloud,
    decode_legacy_pose,
    decode_lio_pcl,
    decode_lio_pose,
)
from k1link.device_plugins.xgrids_k1.viewer.messages import StreamMessage
from k1link.viewer.metrics import BridgeMetrics

POINT_STRUCT = struct.Struct("<fffB3x")
POINT_STRIDE = POINT_STRUCT.size
FRAME_ID = "map"
MAX_TRAJECTORY_POSES = 20_000


@dataclass(frozen=True, slots=True)
class PackedPointCloud:
    data: bytes
    point_count: int


class FoxgloveBridge:
    """Decode verified K1 topics and publish Foxglove-native visualization messages."""

    def __init__(
        self,
        *,
        host: str = "127.0.0.1",
        port: int = 8765,
        metrics: BridgeMetrics | None = None,
    ) -> None:
        self.metrics = metrics or BridgeMetrics()
        self._server = foxglove.start_server(
            name="Mission Core K1 legacy bridge",
            host=host,
            port=port,
            message_backlog_size=32,
        )
        self._points = PointCloudChannel("/k1/points")
        self._pose = PoseInFrameChannel("/k1/pose")
        self._trajectory = SceneUpdateChannel("/k1/trajectory")
        self._metrics = Channel(
            "/k1/metrics",
            schema={
                "type": "object",
                "properties": {
                    "mqtt_to_publish_ms": {"type": ["number", "null"]},
                    "mqtt_to_publish_p50_ms": {"type": ["number", "null"]},
                    "mqtt_to_publish_p95_ms": {"type": ["number", "null"]},
                    "decode_publish_ms": {"type": ["number", "null"]},
                    "point_count": {"type": "integer"},
                    "pcl_fps": {"type": "number"},
                    "pose_fps": {"type": "number"},
                    "preview_dropped": {"type": "integer"},
                },
            },
        )
        self._path: deque[tuple[float, float, float]] = deque(maxlen=MAX_TRAJECTORY_POSES)
        self._last_trajectory_publish_ns = 0
        self._last_point_count = 0
        self._closed = False

    @property
    def port(self) -> int:
        return int(self._server.port)

    @property
    def websocket_url(self) -> str:
        return f"ws://127.0.0.1:{self.port}"

    @property
    def viewer_url(self) -> str:
        return self._server.app_url() or "https://app.foxglove.dev/"

    def process(self, message: StreamMessage) -> None:
        started_ns = time.monotonic_ns()
        self.metrics.received(len(message.payload))
        try:
            if message.topic.endswith("/lio_pcl"):
                self._publish_lio_pcl(decode_lio_pcl(message.payload), message)
            elif message.topic == "RealtimePointcloud":
                self._publish_legacy_pcl(decode_legacy_pointcloud(message.payload), message)
            elif message.topic.endswith("/lio_pose"):
                self._publish_lio_pose(decode_lio_pose(message.payload), message)
            elif message.topic == "RealtimePath":
                self._publish_legacy_pose(decode_legacy_pose(message.payload), message)
            else:
                return
        except StreamDecodeError:
            self.metrics.decode_error()
            return

        published_ns = time.monotonic_ns()
        decode_publish_ms = (published_ns - started_ns) / 1_000_000
        if message.topic.endswith("/lio_pcl") or message.topic == "RealtimePointcloud":
            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))
        if message.received_monotonic_ns is not None:
            self.metrics.record_latency((published_ns - message.received_monotonic_ns) / 1_000_000)
        snapshot = self.metrics.snapshot()
        self._metrics.log(
            {
                "mqtt_to_publish_ms": snapshot["mqtt_to_publish_ms"],
                "mqtt_to_publish_p50_ms": snapshot["mqtt_to_publish_p50_ms"],
                "mqtt_to_publish_p95_ms": snapshot["mqtt_to_publish_p95_ms"],
                "decode_publish_ms": snapshot["decode_publish_ms"],
                "point_count": snapshot["last_point_count"],
                "pcl_fps": snapshot["pcl_fps"],
                "pose_fps": snapshot["pose_fps"],
                "preview_dropped": snapshot["preview_dropped"],
            },
            log_time=message.received_at_epoch_ns,
        )

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        for channel in (self._points, self._pose, self._trajectory, self._metrics):
            channel.close()
        self._server.stop()

    def _publish_lio_pcl(self, frame: LioPointCloudFrame, message: StreamMessage) -> None:
        packed = pack_lio_point_cloud(frame)
        self._publish_point_cloud(packed, message)

    def _publish_legacy_pcl(
        self,
        frame: LegacyPointCloudFrame,
        message: StreamMessage,
    ) -> None:
        packed = pack_legacy_point_cloud(frame)
        self._publish_point_cloud(packed, message)

    def _publish_point_cloud(self, packed: PackedPointCloud, message: StreamMessage) -> None:
        timestamp = _timestamp(message.received_at_epoch_ns)
        self._points.log(
            PointCloud(
                timestamp=timestamp,
                frame_id=FRAME_ID,
                pose=_identity_pose(),
                point_stride=POINT_STRIDE,
                fields=_point_fields(),
                data=packed.data,
            ),
            log_time=message.received_at_epoch_ns,
        )
        self._last_point_count = packed.point_count

    def _publish_lio_pose(self, frame: LioPoseFrame, message: StreamMessage) -> None:
        self._publish_pose(frame.position_xyz, frame.orientation_xyzw, message)

    def _publish_legacy_pose(self, frame: LegacyPoseFrame, message: StreamMessage) -> None:
        self._publish_pose(frame.position_xyz, frame.orientation_xyzw, message)

    def _publish_pose(
        self,
        position_xyz: tuple[float, float, float],
        orientation_xyzw: tuple[float, float, float, float],
        message: StreamMessage,
    ) -> None:
        pose = Pose(
            position=Vector3(x=position_xyz[0], y=position_xyz[1], z=position_xyz[2]),
            orientation=Quaternion(
                x=orientation_xyzw[0],
                y=orientation_xyzw[1],
                z=orientation_xyzw[2],
                w=orientation_xyzw[3],
            ),
        )
        timestamp = _timestamp(message.received_at_epoch_ns)
        self._pose.log(
            PoseInFrame(timestamp=timestamp, frame_id=FRAME_ID, pose=pose),
            log_time=message.received_at_epoch_ns,
        )
        self._path.append(position_xyz)
        now_ns = time.monotonic_ns()
        if (
            len(self._path) > 2
            and len(self._path) % 20
            and now_ns - self._last_trajectory_publish_ns < 200_000_000
        ):
            return
        self._last_trajectory_publish_ns = now_ns
        line_points = [Point3(x=item[0], y=item[1], z=item[2]) for item in self._path]
        self._trajectory.log(
            SceneUpdate(
                entities=[
                    SceneEntity(
                        timestamp=timestamp,
                        frame_id=FRAME_ID,
                        id="k1-trajectory",
                        frame_locked=True,
                        lines=[
                            LinePrimitive(
                                type=LinePrimitiveLineType.LineStrip,
                                thickness=3.0,
                                scale_invariant=True,
                                points=line_points,
                                color=Color(r=0.08, g=0.82, b=1.0, a=1.0),
                            )
                        ],
                    )
                ]
            ),
            log_time=message.received_at_epoch_ns,
        )


def pack_lio_point_cloud(frame: LioPointCloudFrame) -> PackedPointCloud:
    data = bytearray(len(frame.points) * POINT_STRIDE)
    scaler = frame.header.scaler
    for index, point in enumerate(frame.points):
        x, y, z = point.scaled_xyz(scaler)
        POINT_STRUCT.pack_into(data, index * POINT_STRIDE, x, y, z, point.intensity)
    return PackedPointCloud(data=bytes(data), point_count=len(frame.points))


def pack_legacy_point_cloud(frame: LegacyPointCloudFrame) -> PackedPointCloud:
    data = bytearray(len(frame.points) * POINT_STRIDE)
    for index, point in enumerate(frame.points):
        POINT_STRUCT.pack_into(
            data,
            index * POINT_STRIDE,
            point.x,
            point.y,
            point.z,
            point.intensity,
        )
    return PackedPointCloud(data=bytes(data), point_count=len(frame.points))


def _point_fields() -> list[PackedElementField]:
    return [
        PackedElementField(name="x", offset=0, type=PackedElementFieldNumericType.Float32),
        PackedElementField(name="y", offset=4, type=PackedElementFieldNumericType.Float32),
        PackedElementField(name="z", offset=8, type=PackedElementFieldNumericType.Float32),
        PackedElementField(
            name="intensity",
            offset=12,
            type=PackedElementFieldNumericType.Uint8,
        ),
    ]


def _identity_pose() -> Pose:
    return Pose(position=Vector3(), orientation=Quaternion(w=1.0))


def _timestamp(epoch_ns: int) -> Timestamp:
    return Timestamp(epoch_ns // 1_000_000_000, epoch_ns % 1_000_000_000)
