from __future__ import annotations

import math
import statistics
import struct
import threading
import time
from collections import deque
from dataclasses import dataclass
from typing import TypedDict

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.protocol.streams import (
    LegacyPointCloudFrame,
    LegacyPoseFrame,
    LioPointCloudFrame,
    LioPoseFrame,
    StreamDecodeError,
    decode_legacy_pointcloud,
    decode_legacy_pose,
    decode_lio_pcl,
    decode_lio_pose,
)
from k1link.viewer.messages import StreamMessage

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


class MetricsSnapshot(TypedDict):
    messages_received: int
    payload_bytes: int
    pcl_frames: int
    pose_frames: int
    points_published: int
    last_point_count: int
    decode_errors: int
    preview_dropped: int
    pcl_fps: float
    pose_fps: float
    mqtt_to_publish_ms: float | None
    mqtt_to_publish_p50_ms: float | None
    mqtt_to_publish_p95_ms: float | None
    decode_publish_ms: float | None
    trajectory_poses: int


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


class BridgeMetrics:
    """Thread-safe counters shared with the local control API."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._messages_received = 0
        self._payload_bytes = 0
        self._pcl_frames = 0
        self._pose_frames = 0
        self._points_published = 0
        self._last_point_count = 0
        self._decode_errors = 0
        self._preview_dropped = 0
        self._trajectory_poses = 0
        self._pcl_times: deque[int] = deque()
        self._pose_times: deque[int] = deque()
        self._latencies_ms: deque[float] = deque(maxlen=512)
        self._decode_publish_ms: float | None = None

    def received(self, payload_bytes: int) -> None:
        with self._lock:
            self._messages_received += 1
            self._payload_bytes += payload_bytes

    def published_pcl(self, point_count: int, now_ns: int, decode_publish_ms: float) -> None:
        with self._lock:
            self._pcl_frames += 1
            self._points_published += point_count
            self._last_point_count = point_count
            self._decode_publish_ms = decode_publish_ms
            self._pcl_times.append(now_ns)
            _trim_rate_window(self._pcl_times, now_ns)

    def published_pose(self, now_ns: int, decode_publish_ms: float, path_size: int) -> None:
        with self._lock:
            self._pose_frames += 1
            self._decode_publish_ms = decode_publish_ms
            self._trajectory_poses = path_size
            self._pose_times.append(now_ns)
            _trim_rate_window(self._pose_times, now_ns)

    def record_latency(self, milliseconds: float) -> None:
        if not math.isfinite(milliseconds) or milliseconds < 0:
            return
        with self._lock:
            self._latencies_ms.append(milliseconds)

    def decode_error(self) -> None:
        with self._lock:
            self._decode_errors += 1

    def preview_dropped(self) -> None:
        with self._lock:
            self._preview_dropped += 1

    def snapshot(self) -> MetricsSnapshot:
        now_ns = time.monotonic_ns()
        with self._lock:
            _trim_rate_window(self._pcl_times, now_ns)
            _trim_rate_window(self._pose_times, now_ns)
            latencies = list(self._latencies_ms)
            last_latency = latencies[-1] if latencies else None
            p50 = statistics.median(latencies) if latencies else None
            p95 = _percentile(latencies, 0.95) if latencies else None
            return {
                "messages_received": self._messages_received,
                "payload_bytes": self._payload_bytes,
                "pcl_frames": self._pcl_frames,
                "pose_frames": self._pose_frames,
                "points_published": self._points_published,
                "last_point_count": self._last_point_count,
                "decode_errors": self._decode_errors,
                "preview_dropped": self._preview_dropped,
                "pcl_fps": _window_rate(self._pcl_times),
                "pose_fps": _window_rate(self._pose_times),
                "mqtt_to_publish_ms": _rounded(last_latency),
                "mqtt_to_publish_p50_ms": _rounded(p50),
                "mqtt_to_publish_p95_ms": _rounded(p95),
                "decode_publish_ms": _rounded(self._decode_publish_ms),
                "trajectory_poses": self._trajectory_poses,
            }


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)


def _trim_rate_window(values: deque[int], now_ns: int) -> None:
    cutoff = now_ns - 1_000_000_000
    while values and values[0] < cutoff:
        values.popleft()


def _window_rate(values: deque[int]) -> float:
    if len(values) < 2:
        return float(len(values))
    elapsed = (values[-1] - values[0]) / 1_000_000_000
    return len(values) / max(elapsed, 1.0)


def _percentile(values: list[float], fraction: float) -> float:
    if not values:
        raise ValueError("cannot calculate a percentile of an empty sample")
    ordered = sorted(values)
    index = min(len(ordered) - 1, math.ceil(len(ordered) * fraction) - 1)
    return ordered[index]


def _rounded(value: float | None) -> float | None:
    return None if value is None else round(value, 3)