"""WebSocket LED client — broadcasts pixel data to connected WebSocket clients."""

import asyncio
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from wled_controller.core.devices.led_client import DeviceHealth, LEDClient
from wled_controller.utils import get_logger

logger = get_logger(__name__)


class WSDeviceBroadcaster:
    """Global registry of WebSocket clients subscribed to WS device streams.

    Each WS device (identified by device_id) can have zero or more connected
    WebSocket clients.  The WSLEDClient.send_pixels() method uses this to
    broadcast frames.
    """

    def __init__(self):
        self._clients: Dict[str, List] = {}

    def add_client(self, device_id: str, ws) -> None:
        self._clients.setdefault(device_id, []).append(ws)
        logger.info(
            "WS device %s: client connected (%d total)",
            device_id, len(self._clients[device_id]),
        )

    def remove_client(self, device_id: str, ws) -> None:
        clients = self._clients.get(device_id)
        if clients and ws in clients:
            clients.remove(ws)
            logger.info(
                "WS device %s: client disconnected (%d remaining)",
                device_id, len(clients),
            )

    def get_clients(self, device_id: str) -> List:
        return self._clients.get(device_id, [])


_broadcaster = WSDeviceBroadcaster()


def get_ws_broadcaster() -> WSDeviceBroadcaster:
    return _broadcaster


def parse_ws_url(url: str) -> str:
    """Extract device_id from a ws:// URL.

    Format: ws://device-id
    """
    if url.startswith("ws://"):
        return url[5:]
    return url


class WSLEDClient(LEDClient):
    """Broadcasts binary pixel data to WebSocket clients via the global broadcaster."""

    def __init__(self, url: str, led_count: int = 0, **kwargs):
        self._device_id = parse_ws_url(url)
        self._led_count = led_count
        self._connected = False

    async def connect(self) -> bool:
        self._connected = True
        logger.info("WS device client connected for device %s", self._device_id)
        return True

    async def close(self) -> None:
        self._connected = False
        logger.info("WS device client closed for device %s", self._device_id)

    @property
    def is_connected(self) -> bool:
        return self._connected

    async def send_pixels(
        self,
        pixels: Union[List[Tuple[int, int, int]], np.ndarray],
        brightness: int = 255,
    ) -> bool:
        if not self._connected:
            return False

        clients = _broadcaster.get_clients(self._device_id)
        if not clients:
            return True

        # Build binary frame: [brightness_byte][R G B R G B ...]
        if isinstance(pixels, np.ndarray):
            pixel_bytes = pixels.astype(np.uint8).tobytes()
        else:
            pixel_bytes = bytes(c for rgb in pixels for c in rgb)

        data = bytes([brightness]) + pixel_bytes

        async def _send_safe(ws):
            try:
                await ws.send_bytes(data)
                return True
            except Exception:
                return False

        results = await asyncio.gather(*[_send_safe(ws) for ws in clients])

        disconnected = [ws for ws, ok in zip(clients, results) if not ok]
        for ws in disconnected:
            _broadcaster.remove_client(self._device_id, ws)

        return True

    @classmethod
    async def check_health(
        cls,
        url: str,
        http_client,
        prev_health: Optional[DeviceHealth] = None,
    ) -> DeviceHealth:
        return DeviceHealth(
            online=True,
            latency_ms=0.0,
            last_checked=datetime.now(timezone.utc),
        )
