from __future__ import annotations

import asyncio
import json
import time
from typing import Any

from fastapi import WebSocket, WebSocketDisconnect

from server.telemetry.listener import telemetry_listener
from server.telemetry.data import TelemetryData
from utils.logger import get_logger

logger = get_logger(__name__)


class ConnectionManager:
    def __init__(self):
        self._connections: dict[str, WebSocket] = {}
        self._counter = 0
        self._broadcast_task: asyncio.Task | None = None
        self._pending_data: TelemetryData | None = None

    async def connect(self, ws: WebSocket) -> str:
        await ws.accept()
        self._counter += 1
        cid = f"client_{self._counter}"
        self._connections[cid] = ws
        logger.info("WS client connected: %s (total: %d)", cid, len(self._connections))
        if not self._broadcast_task or self._broadcast_task.done():
            self._broadcast_task = asyncio.create_task(self._broadcast_loop())
        return cid

    def disconnect(self, cid: str):
        self._connections.pop(cid, None)
        logger.info("WS client disconnected: %s (total: %d)", cid, len(self._connections))
        if not self._connections and self._broadcast_task:
            self._broadcast_task.cancel()
            self._broadcast_task = None

    def push_telemetry(self, data: TelemetryData):
        self._pending_data = data

    async def broadcast(self, message: dict[str, Any]):
        dead = []
        for cid, ws in self._connections.items():
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(cid)
        for cid in dead:
            self.disconnect(cid)

    async def _broadcast_loop(self):
        last_sent = 0.0
        throttle_interval = 1.0 / 30.0
        try:
            while self._connections:
                now = time.time()
                if now - last_sent >= throttle_interval and self._pending_data:
                    data = self._pending_data
                    self._pending_data = None
                    last_sent = now
                    await self.broadcast({
                        "type": "telemetry",
                        "data": data.to_dict(),
                    })
                await asyncio.sleep(0.01)
        except asyncio.CancelledError:
            pass

    @property
    def client_count(self) -> int:
        return len(self._connections)


ws_manager = ConnectionManager()


def on_telemetry(data: TelemetryData):
    ws_manager.push_telemetry(data)


telemetry_listener.on_data(on_telemetry)