import struct

# Protocol type bytes (first byte of UDP packet)
PROTO_DRGB = 0x02
PROTO_DRGBW = 0x03
PROTO_DNRGB = 0x04
PROTO_ANIM_CMD = 0xAC

ANIMATION_NAMES = {
    0: 'solid_color', 1: 'chase', 2: 'pulse',
    3: 'rainbow', 4: 'strobe', 5: 'color_wipe', 6: 'fire'
}


def classify_packet(data):
    """Return packet type string from first byte, or None if unknown."""
    if not data:
        return None
    b = data[0]
    if b == PROTO_DRGB:
        return 'drgb'
    elif b == PROTO_DRGBW:
        return 'drgbw'
    elif b == PROTO_DNRGB:
        return 'dnrgb'
    elif b == PROTO_ANIM_CMD:
        return 'anim_cmd'
    return None


def parse_drgb(data):
    """Parse DRGB packet -> list of (R,G,B) tuples.

    Format: [0x02][timeout][R0][G0][B0]...[Rn][Gn][Bn]
    """
    n = (len(data) - 2) // 3
    pixels = []
    for i in range(n):
        off = 2 + i * 3
        pixels.append((data[off], data[off + 1], data[off + 2]))
    return pixels


def parse_drgbw(data):
    """Parse DRGBW packet -> list of (R,G,B,W) tuples.

    Format: [0x03][timeout][R0][G0][B0][W0]...[Rn][Gn][Bn][Wn]
    """
    n = (len(data) - 2) // 4
    pixels = []
    for i in range(n):
        off = 2 + i * 4
        pixels.append((data[off], data[off + 1], data[off + 2], data[off + 3]))
    return pixels


def parse_dnrgb(data):
    """Parse DNRGB packet -> (start_index, list of (R,G,B) tuples).

    Format: [0x04][timeout][start_hi][start_lo][R0][G0][B0]...
    Start index is big-endian uint16.
    """
    start_index = struct.unpack('>H', data[2:4])[0]
    n = (len(data) - 4) // 3
    pixels = []
    for i in range(n):
        off = 4 + i * 3
        pixels.append((data[off], data[off + 1], data[off + 2]))
    return start_index, pixels


def parse_animation_cmd(data):
    """Parse animation command packet -> dict with animation name and params.

    Format (16 bytes, big-endian):
      Byte 0:     0xAC marker
      Byte 1:     version (1)
      Byte 2:     animation type ID (0-6)
      Byte 3:     flags (bit 0 = reverse)
      Bytes 4-7:  speed (float32 big-endian)
      Bytes 8-10: primary color R, G, B
      Bytes 11-13: background color R, G, B
      Byte 14:    white channel
      Byte 15:    density/size/cooling
    """
    if len(data) < 16 or data[0] != PROTO_ANIM_CMD:
        return None
    anim_id = data[2]
    flags = data[3]
    speed = struct.unpack('>f', data[4:8])[0]
    anim_name = ANIMATION_NAMES.get(anim_id)
    if anim_name is None:
        return None
    return {
        'animation': anim_name,
        'speed': speed,
        'color': (data[8], data[9], data[10]),
        'bg_color': (data[11], data[12], data[13]),
        'white': data[14],
        'density': data[15],
        'reverse': bool(flags & 0x01),
    }


class FrameAssembler:
    """Reassembles multi-packet DNRGB frames.

    For 300-LED strips, frames fit in one packet so reassembly never triggers.
    Implemented for completeness and future strip sizes.
    """

    def __init__(self, expected_count):
        self.expected_count = expected_count
        self._buf = {}  # start_index -> list of (R,G,B)
        self._last_ms = 0
        self._TIMEOUT_MS = 200

    def feed(self, start_index, pixels, now_ms):
        """Feed a DNRGB chunk. Returns complete frame list or None."""
        # Timeout: discard stale partial frames
        if self._buf and (now_ms - self._last_ms) > self._TIMEOUT_MS:
            self._buf.clear()

        self._last_ms = now_ms

        for i, px in enumerate(pixels):
            self._buf[start_index + i] = px

        if len(self._buf) >= self.expected_count:
            frame = [self._buf.get(i, (0, 0, 0)) for i in range(self.expected_count)]
            self._buf.clear()
            return frame

        return None

    def reset(self):
        self._buf.clear()