"""DDP (Distributed Display Protocol) client for WLED."""

import asyncio
import struct
from dataclasses import dataclass
from typing import Dict, List, Tuple

import numpy as np

from ledgrab.utils import get_logger

logger = get_logger(__name__)

# DDP color order codes (matches WLED's order enum)
COLOR_ORDER_MAP: Dict[int, Tuple[int, int, int]] = {
    0: (1, 0, 2),  # GRB → indices into (R,G,B)
    1: (0, 1, 2),  # RGB → no reorder
    2: (2, 0, 1),  # BRG
    3: (0, 2, 1),  # RBG
    4: (2, 1, 0),  # BGR
    5: (1, 2, 0),  # GBR
}


@dataclass
class BusConfig:
    """Physical LED bus/output configuration from WLED."""

    start: int  # First LED index
    length: int  # Number of LEDs on this bus
    color_order: int  # Color order code (0=GRB, 1=RGB, etc.)


class DDPClient:
    """UDP DDP client for sending pixel data to WLED devices."""

    DDP_PORT = 4048
    DDP_FLAGS_VER1 = 0x40  # VER=1, all other flags 0
    DDP_FLAGS_PUSH = 0x01  # PUSH flag (set on last packet of a frame)
    DDP_TYPE_RGB = 0x01

    # Pre-built struct.Struct for the 10-byte DDP header (avoids per-call format parsing)
    _HEADER_STRUCT = struct.Struct("!BBB B I H")

    def __init__(self, host: str, port: int = DDP_PORT, rgbw: bool = False):
        """Initialize DDP client.

        Args:
            host: WLED device IP address or hostname
            port: DDP port (default 4048)
            rgbw: True for RGBW LEDs (4 bytes/LED), False for RGB (3 bytes/LED)
        """
        self.host = host
        self.port = port
        self.rgbw = rgbw
        self._transport = None
        self._protocol = None
        self._sequence = 0
        self._buses: List[BusConfig] = []
        # Pre-allocated RGBW buffer (resized on demand)
        self._rgbw_buf: np.ndarray | None = None
        self._rgbw_buf_n: int = 0
        # Pre-allocated send buffer (header + payload). Sized lazily on first
        # send so we never allocate fresh bytes per frame on the hot path.
        self._send_buf: bytearray | None = None
        self._send_view: memoryview | None = None

    async def connect(self):
        """Establish UDP connection."""
        try:
            loop = asyncio.get_running_loop()
            self._transport, self._protocol = await loop.create_datagram_endpoint(
                asyncio.DatagramProtocol, remote_addr=(self.host, self.port)
            )
            logger.info(f"DDP client connected to {self.host}:{self.port}")
            return True
        except Exception as e:
            logger.error(f"Failed to connect DDP client: {e}")
            raise

    async def close(self):
        """Close the connection."""
        if self._transport:
            self._transport.close()
            self._transport = None
            self._protocol = None
        logger.debug(f"Closed DDP connection to {self.host}:{self.port}")

    def set_buses(self, buses: List[BusConfig]) -> None:
        """Set WLED bus configurations for per-bus color order reordering.

        Args:
            buses: List of BusConfig from WLED device
        """
        self._buses = buses
        for bus in buses:
            order_name = {0: "GRB", 1: "RGB", 2: "BRG", 3: "RBG", 4: "BGR", 5: "GBR"}
            logger.info(
                f"DDP bus: LEDs {bus.start}-{bus.start + bus.length - 1}, "
                f"color order={order_name.get(bus.color_order, '?')} ({bus.color_order})"
            )

    def _ensure_send_buf(self, capacity: int) -> None:
        """Lazily allocate / grow the per-instance send buffer.

        ``capacity`` is the largest packet we may emit (header + payload).
        Once sized, the buffer is reused for every subsequent send so the
        hot path stays allocation-free.
        """
        buf = self._send_buf
        if buf is None or len(buf) < capacity:
            self._send_buf = bytearray(capacity)
            self._send_view = memoryview(self._send_buf)

    def _emit_packet(
        self,
        payload: memoryview,
        offset: int,
        sequence: int,
        push: bool,
    ) -> None:
        """Pack header + payload into the pre-allocated send buffer and emit.

        DDP packet layout (10-byte header):
            [0]   Flags (VER1 | PUSH on last)
            [1]   Sequence
            [2]   Data type (0x01 = RGB)
            [3]   Source/Destination ID
            [4-7] Data offset (big-endian)
            [8-9] Data length (big-endian)
            [10+] Pixel data
        """
        flags = self.DDP_FLAGS_VER1 | (self.DDP_FLAGS_PUSH if push else 0)
        data_len = len(payload)
        self._ensure_send_buf(10 + data_len)
        buf = self._send_buf
        view = self._send_view
        # Fill header into pre-allocated buffer (no allocation)
        self._HEADER_STRUCT.pack_into(
            buf, 0, flags, sequence, self.DDP_TYPE_RGB, 0x01, offset, data_len
        )
        # Copy payload bytes into buffer (single memcpy)
        view[10 : 10 + data_len] = payload
        # asyncio's selector_datagram_transport.sendto fast-path calls
        # socket.sendto(data) which accepts a buffer-like; it only copies to
        # bytes when the OS send buffer is full and the datagram must be
        # queued. So passing a memoryview is safe and avoids `bytes(...)`.
        self._transport.sendto(view[: 10 + data_len])

    def _reorder_pixels_numpy(self, pixel_array: np.ndarray) -> np.ndarray:
        """Apply per-bus color order reordering using numpy fancy indexing.

        WLED may not apply per-bus color order conversion for DDP data on
        all buses (observed in multi-bus setups). We reorder pixel channels
        here so the hardware receives the correct byte order directly.

        Args:
            pixel_array: (N, 3) uint8 numpy array in RGB order

        Returns:
            Reordered array (may be a view or copy depending on buses)
        """
        if not self._buses:
            return pixel_array

        result = pixel_array.copy()
        for bus in self._buses:
            order_map = COLOR_ORDER_MAP.get(bus.color_order)
            if not order_map or order_map == (0, 1, 2):
                continue  # RGB order = no reordering needed

            start = bus.start
            end = min(bus.start + bus.length, len(result))
            result[start:end] = result[start:end][:, order_map]

        return result

    def _send_buffer(self, payload_view: memoryview, bpp: int, max_packet_size: int) -> int:
        """Chunk and emit a contiguous payload via DDP. Returns packet count.

        ``payload_view`` is a 1-D bytes-like view; the caller guarantees its
        length is a multiple of ``bpp``.  Each emitted packet is sized to a
        whole number of pixels so RGB channels never split across packets.
        """
        total_bytes = len(payload_view)
        max_payload = max_packet_size - 10  # 10-byte header
        bytes_per_packet = (max_payload // bpp) * bpp
        if bytes_per_packet <= 0:
            bytes_per_packet = bpp  # degenerate guard

        num_packets = (total_bytes + bytes_per_packet - 1) // bytes_per_packet
        for i in range(num_packets):
            start = i * bytes_per_packet
            end = total_bytes if (i == num_packets - 1) else (start + bytes_per_packet)
            self._sequence = (self._sequence + 1) % 256
            self._emit_packet(
                payload_view[start:end],
                offset=start,
                sequence=self._sequence,
                push=(i == num_packets - 1),
            )
        return num_packets

    async def send_pixels(
        self, pixels: List[Tuple[int, int, int]], max_packet_size: int = 1400
    ) -> bool:
        """Send pixel data via DDP.

        Args:
            pixels: List of (R, G, B) tuples or an (N, 3) uint8 numpy array
            max_packet_size: Maximum UDP packet size (default 1400 bytes for safety)

        Returns:
            True if successful

        Raises:
            RuntimeError: If send fails
        """
        if not self._transport:
            raise RuntimeError("DDP client not connected")

        try:
            # Accept numpy arrays directly to avoid per-pixel Python loop
            if isinstance(pixels, np.ndarray):
                pixel_array = pixels
            else:
                pixel_array = np.array(pixels, dtype=np.uint8)

            self.send_pixels_numpy(pixel_array, max_packet_size=max_packet_size)
            return True

        except Exception as e:
            logger.error("Failed to send DDP pixels: %s", e)
            raise RuntimeError(f"DDP send failed: {e}")

    def send_pixels_numpy(self, pixel_array: np.ndarray, max_packet_size: int = 1400) -> bool:
        """Send pixel data via DDP from a numpy array — no per-pixel Python loops.

        Args:
            pixel_array: (N, 3) uint8 numpy array of RGB values
            max_packet_size: Maximum UDP packet size (default 1400 bytes for safety)

        Returns:
            True if successful
        """
        if not self._transport:
            raise RuntimeError("DDP client not connected")

        # Handle RGBW: copy RGB into pre-allocated (N, 4) buffer
        if self.rgbw:
            n = pixel_array.shape[0]
            if n != self._rgbw_buf_n:
                self._rgbw_buf = np.zeros((n, 4), dtype=np.uint8)
                self._rgbw_buf_n = n
            self._rgbw_buf[:, :3] = pixel_array
            pixel_array = self._rgbw_buf

        bpp = 4 if self.rgbw else 3
        # Get a 1-D bytes view of the pixel buffer with no allocation when
        # the array is already C-contiguous (the common case).
        if not pixel_array.flags["C_CONTIGUOUS"]:
            pixel_array = np.ascontiguousarray(pixel_array)
        # ``cast('B')`` on a memoryview of a numpy array returns a 1-D byte
        # view; total length == nbytes.
        payload_view = memoryview(pixel_array).cast("B")
        self._send_buffer(payload_view, bpp, max_packet_size)
        return True

    async def __aenter__(self):
        """Async context manager entry."""
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        await self.close()
