"""DDP (Distributed Display Protocol) client for WLED."""
import asyncio
import struct
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import numpy as np

from wled_controller.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

    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: Optional[np.ndarray] = None
        self._rgbw_buf_n: int = 0

    async def connect(self):
        """Establish UDP connection."""
        try:
            loop = asyncio.get_event_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 _build_ddp_packet(
        self,
        rgb_data: bytes,
        offset: int = 0,
        sequence: int = 1,
        push: bool = False,
    ) -> bytes:
        """Build a DDP packet.

        DDP packet format (10-byte header + data):
        - Byte 0: Flags (VER1 | PUSH on last packet)
        - Byte 1: Sequence number
        - Byte 2: Data type (0x01 = RGB)
        - Byte 3: Source/Destination ID
        - Bytes 4-7: Data offset (4 bytes, big-endian)
        - Bytes 8-9: Data length (2 bytes, big-endian)
        - Bytes 10+: Pixel data

        Args:
            rgb_data: RGB pixel data as bytes
            offset: Byte offset (pixel_index * 3)
            sequence: Sequence number (0-255)
            push: True for the last packet of a frame

        Returns:
            Complete DDP packet as bytes
        """
        flags = self.DDP_FLAGS_VER1
        if push:
            flags |= self.DDP_FLAGS_PUSH
        data_type = self.DDP_TYPE_RGB
        source_id = 0x01
        data_len = len(rgb_data)

        # Build header (10 bytes)
        header = struct.pack(
            '!BBB B I H',  # Network byte order (big-endian)
            flags,         # Flags
            sequence,      # Sequence
            data_type,     # Data type
            source_id,     # Source/Destination
            offset,        # Data offset (4 bytes)
            data_len       # Data length (2 bytes)
        )

        return header + rgb_data

    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

    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
            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:
            # Send plain RGB — WLED handles per-bus color order conversion
            # internally when outputting to hardware.
            # Accept numpy arrays directly to avoid per-pixel Python loop
            bpp = 4 if self.rgbw else 3  # bytes per pixel
            if isinstance(pixels, np.ndarray):
                pixel_array = pixels
            else:
                pixel_array = np.array(pixels, dtype=np.uint8)
            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
            pixel_bytes = pixel_array.tobytes()

            total_bytes = len(pixel_bytes)
            # Align payload to full pixels (multiple of bpp) to avoid splitting
            # a pixel's channels across packets
            max_payload = max_packet_size - 10  # 10-byte header
            bytes_per_packet = (max_payload // bpp) * bpp

            # Split into multiple packets if needed
            num_packets = (total_bytes + bytes_per_packet - 1) // bytes_per_packet

            logger.debug(
                f"DDP: Sending {len(pixels)} pixels ({total_bytes} bytes) "
                f"in {num_packets} packet(s) to {self.host}:{self.port}"
            )

            for i in range(num_packets):
                start = i * bytes_per_packet
                end = min(start + bytes_per_packet, total_bytes)
                chunk = pixel_bytes[start:end]
                is_last = (i == num_packets - 1)

                # Increment sequence number
                self._sequence = (self._sequence + 1) % 256

                # Set PUSH flag on the last packet to signal frame completion
                packet = self._build_ddp_packet(
                    chunk, offset=start,
                    sequence=self._sequence, push=is_last,
                )
                self._transport.sendto(packet)

                logger.debug(
                    f"Sent DDP packet {i+1}/{num_packets}: "
                    f"{len(chunk)} bytes at offset {start}"
                    f"{' [PUSH]' if is_last else ''}"
                )

            return True

        except Exception as e:
            logger.error(f"Failed to send DDP pixels: {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

        pixel_bytes = pixel_array.tobytes()

        bpp = 4 if self.rgbw else 3
        total_bytes = len(pixel_bytes)
        max_payload = max_packet_size - 10  # 10-byte header
        bytes_per_packet = (max_payload // bpp) * bpp
        num_packets = (total_bytes + bytes_per_packet - 1) // bytes_per_packet

        for i in range(num_packets):
            start = i * bytes_per_packet
            end = min(start + bytes_per_packet, total_bytes)
            chunk = pixel_bytes[start:end]
            is_last = (i == num_packets - 1)
            self._sequence = (self._sequence + 1) % 256
            packet = self._build_ddp_packet(
                chunk, offset=start,
                sequence=self._sequence, push=is_last,
            )
            self._transport.sendto(packet)

        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()
