"""LIFX LAN LED client.

LIFX bulbs and lightstrips accept a binary UDP protocol on port 56700.
Every packet has a 36-byte header (frame + frame-address + protocol-header)
followed by a type-specific payload. Colors are HSBK 16-bit per channel.

URL scheme: ``lifx://<host>[:port]`` or bare ``<host>[:port]``. Default port 56700.

LIFX bulbs are reachable two ways:
  * Unicast — set the ``target`` field to the bulb's 48-bit MAC.
  * Broadcast — set ``target`` to all zeros and ``tagged=1``; all bulbs on
    the subnet act on the message. We use this for the SetColor hot path
    so we don't have to learn the MAC of every device a user owns.

Reference: https://lan.developer.lifx.com/docs/header-description
"""

from __future__ import annotations

import asyncio
import socket
import struct
import time
from datetime import datetime, timezone
from typing import List, Tuple
from urllib.parse import urlparse

import numpy as np

from ledgrab.core.devices.led_client import DeviceHealth, LEDClient
from ledgrab.core.devices.pixel_reduce import average_color as _average_color
from ledgrab.utils import get_logger

logger = get_logger(__name__)

LIFX_PORT = 56700
DEFAULT_MIN_INTERVAL_S = 0.05  # ~20 Hz — LIFX rate-limit guidance is 20/sec

# Message types we care about
MSG_GET_SERVICE = 2
MSG_STATE_SERVICE = 3
MSG_SET_POWER = 21
MSG_SET_COLOR = 102

# Frame field byte 0 of the protocol header: tagged=1, addressable=1, protocol=1024
_FRAME_TAGGED = 0x3400
_FRAME_UNTAGGED = 0x1400

_SOURCE_ID = 0x4C474752  # "LGGR" — identifies LedGrab in protocol logs


def parse_lifx_url(url: str) -> Tuple[str, int]:
    """Pull ``(host, port)`` from ``lifx://host[:port]`` or bare ``host[:port]``."""
    if not url:
        raise ValueError("LIFX URL is empty")
    raw = url.strip()
    if "://" in raw:
        parsed = urlparse(raw)
        host = parsed.hostname or ""
        port = parsed.port or LIFX_PORT
    else:
        parsed = urlparse(f"lifx://{raw}")
        host = parsed.hostname or ""
        port = parsed.port or LIFX_PORT
    if not host:
        raise ValueError(f"LIFX URL has no host: {url!r}")
    return host, port


def rgb_to_hsbk(r: int, g: int, b: int, kelvin: int = 3500) -> Tuple[int, int, int, int]:
    """Convert 8-bit RGB to LIFX 16-bit HSBK.

    The ``kelvin`` channel is irrelevant when saturation > 0 (the bulb
    interprets it as a hint, not a hard temperature), so we leave it at the
    LIFX default of ~3500 K. Outputs are clamped uint16.
    """
    r_n = max(0, min(255, r)) / 255.0
    g_n = max(0, min(255, g)) / 255.0
    b_n = max(0, min(255, b)) / 255.0
    c_max = max(r_n, g_n, b_n)
    c_min = min(r_n, g_n, b_n)
    delta = c_max - c_min

    # Hue (0-360 degrees → 0-65535)
    if delta == 0:
        h = 0.0
    elif c_max == r_n:
        h = 60.0 * (((g_n - b_n) / delta) % 6)
    elif c_max == g_n:
        h = 60.0 * (((b_n - r_n) / delta) + 2)
    else:
        h = 60.0 * (((r_n - g_n) / delta) + 4)
    hue_u16 = int((h / 360.0) * 65535) & 0xFFFF

    # Saturation (0-1 → 0-65535)
    sat_u16 = 0 if c_max == 0 else int((delta / c_max) * 65535) & 0xFFFF
    # Brightness (0-1 → 0-65535)
    bri_u16 = int(c_max * 65535) & 0xFFFF
    kelvin_u16 = max(2500, min(9000, kelvin)) & 0xFFFF
    return hue_u16, sat_u16, bri_u16, kelvin_u16


def _build_packet(
    *,
    msg_type: int,
    payload: bytes,
    target_mac: bytes = b"\x00\x00\x00\x00\x00\x00",
    sequence: int = 0,
    res_required: bool = False,
    ack_required: bool = False,
    tagged: bool = False,
) -> bytes:
    """Pack a LIFX packet: 36-byte header + payload.

    See https://lan.developer.lifx.com/docs/header-description for the
    bit-level field layout. We construct the three sub-headers separately
    so the magic numbers are scoped to the fields they belong to.
    """
    size = 36 + len(payload)
    # Frame header (8 bytes): size(2) | origin/tagged/addressable/protocol(2) | source(4)
    frame_field = _FRAME_TAGGED if tagged else _FRAME_UNTAGGED
    frame = struct.pack("<HHI", size, frame_field, _SOURCE_ID)
    # Frame address (16 bytes): target(8) | reserved(6) | res/ack flags(1) | sequence(1)
    flags = (0x01 if res_required else 0) | (0x02 if ack_required else 0)
    frame_addr = (
        target_mac + b"\x00\x00" + b"\x00\x00\x00\x00\x00\x00" + bytes([flags, sequence & 0xFF])
    )
    # Protocol header (12 bytes): reserved(8) | type(2) | reserved(2)
    proto = b"\x00" * 8 + struct.pack("<HH", msg_type, 0)
    return frame + frame_addr + proto + payload


def _build_set_color_payload(h: int, s: int, b: int, k: int, duration_ms: int = 0) -> bytes:
    """SetColor payload: reserved(1) | HSBK(8) | duration(4)."""
    return b"\x00" + struct.pack(
        "<HHHHI", h & 0xFFFF, s & 0xFFFF, b & 0xFFFF, k & 0xFFFF, duration_ms & 0xFFFFFFFF
    )


def _build_set_power_payload(on: bool, duration_ms: int = 0) -> bytes:
    """SetPower payload: level(2) | duration(4). Level is 0 or 65535."""
    return struct.pack("<HI", 65535 if on else 0, duration_ms & 0xFFFFFFFF)


def _parse_state_service_reply(raw: bytes) -> dict | None:
    """Parse a LIFX StateService (discovery) reply.

    Returns ``{"mac": "aabbccddeeff", "service": int, "port": int}`` or
    ``None`` if the payload isn't a StateService.
    """
    if len(raw) < 36 + 5:
        return None
    # Read msg_type from the protocol header at offset 32
    msg_type = struct.unpack_from("<H", raw, 32)[0]
    if msg_type != MSG_STATE_SERVICE:
        return None
    # Target MAC at offset 8, 6 bytes are the MAC.
    mac_bytes = raw[8:14]
    mac = mac_bytes.hex()
    # Payload at offset 36: service(1) + port(4)
    service, port = struct.unpack_from("<BI", raw, 36)
    return {"mac": mac, "service": int(service), "port": int(port)}


class _LIFXProtocol(asyncio.DatagramProtocol):
    """Write-only datagram protocol. Inbound replies dropped silently."""

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        # LIFX bulbs sometimes echo back state on broadcast. We don't need it
        # for streaming ambilight — discard.
        pass

    def error_received(self, exc):
        logger.debug("LIFX UDP error: %s", exc)


class LIFXClient(LEDClient):
    """LEDClient for a single LIFX bulb on the LAN."""

    def __init__(
        self,
        url: str,
        led_count: int = 1,
        *,
        min_interval_s: float = DEFAULT_MIN_INTERVAL_S,
    ):
        host, port = parse_lifx_url(url)
        self._host = host
        self._port = port
        self._led_count = led_count
        self._min_interval_s = max(0.0, min_interval_s)
        self._transport: asyncio.DatagramTransport | None = None
        self._protocol: _LIFXProtocol | None = None
        self._connected = False
        self._next_tx_at: float = 0.0
        self._sequence: int = 0

    @property
    def host(self) -> str:
        return self._host

    @property
    def port(self) -> int:
        return self._port

    @property
    def is_connected(self) -> bool:
        return self._connected and self._transport is not None

    @property
    def device_led_count(self) -> int | None:
        return self._led_count or None

    async def connect(self) -> bool:
        if self._connected and self._transport is not None:
            return True
        loop = asyncio.get_running_loop()
        try:
            transport, protocol = await loop.create_datagram_endpoint(
                _LIFXProtocol, remote_addr=(self._host, self._port)
            )
        except OSError as exc:
            raise RuntimeError(f"Failed to open UDP to LIFX at {self._host}: {exc}") from exc
        self._transport = transport
        self._protocol = protocol  # type: ignore[assignment]
        self._connected = True
        logger.info("LIFXClient connected to %s:%d", self._host, self._port)
        return True

    async def close(self) -> None:
        if self._transport is not None:
            try:
                self._transport.close()
            except OSError:
                pass
        self._transport = None
        self._protocol = None
        self._connected = False

    def _next_sequence(self) -> int:
        self._sequence = (self._sequence + 1) & 0xFF
        return self._sequence

    def _send(self, msg_type: int, payload: bytes) -> None:
        """Build and send one LIFX packet. Caller must hold an open transport."""
        assert self._transport is not None
        packet = _build_packet(
            msg_type=msg_type,
            payload=payload,
            sequence=self._next_sequence(),
            tagged=True,  # broadcast within the unicast UDP socket — bulb still acts on it
        )
        self._transport.sendto(packet)

    async def send_pixels(
        self,
        pixels: List[Tuple[int, int, int]] | np.ndarray,
        brightness: int = 255,
    ) -> bool:
        """Average the strip → HSBK → SetColor."""
        if not self.is_connected:
            raise RuntimeError("LIFXClient not connected")
        now = time.monotonic()
        if now < self._next_tx_at:
            return True
        r, g, b = _average_color(pixels)
        if brightness < 255:
            scale = max(0, min(255, brightness)) / 255.0
            r = int(r * scale)
            g = int(g * scale)
            b = int(b * scale)
        h, s, br, k = rgb_to_hsbk(r, g, b)
        self._send(MSG_SET_COLOR, _build_set_color_payload(h, s, br, k, duration_ms=0))
        self._next_tx_at = now + self._min_interval_s
        return True

    def send_pixels_fast(
        self,
        pixels: List[Tuple[int, int, int]] | np.ndarray,
        brightness: int = 255,
    ) -> None:
        """Synchronous variant — same shape, runs on the hot loop."""
        if not self.is_connected or self._transport is None:
            raise RuntimeError("LIFXClient not connected")
        now = time.monotonic()
        if now < self._next_tx_at:
            return
        r, g, b = _average_color(pixels)
        if brightness < 255:
            scale = max(0, min(255, brightness)) / 255.0
            r = int(r * scale)
            g = int(g * scale)
            b = int(b * scale)
        h, s, br, k = rgb_to_hsbk(r, g, b)
        self._send(MSG_SET_COLOR, _build_set_color_payload(h, s, br, k, duration_ms=0))
        self._next_tx_at = now + self._min_interval_s

    @property
    def supports_fast_send(self) -> bool:
        return True

    async def set_color(self, r: int, g: int, b: int) -> None:
        if not self.is_connected:
            raise RuntimeError("LIFXClient not connected")
        h, s, br, k = rgb_to_hsbk(r, g, b)
        self._send(MSG_SET_COLOR, _build_set_color_payload(h, s, br, k, duration_ms=0))

    async def set_power(self, on: bool) -> None:
        if not self.is_connected:
            raise RuntimeError("LIFXClient not connected")
        self._send(MSG_SET_POWER, _build_set_power_payload(on))

    @classmethod
    async def check_health(
        cls,
        url: str,
        http_client,
        prev_health: DeviceHealth | None = None,
    ) -> DeviceHealth:
        """Send GetService and wait briefly for a StateService reply."""
        now = datetime.now(timezone.utc)
        try:
            host, port = parse_lifx_url(url)
        except ValueError as exc:
            return DeviceHealth(online=False, last_checked=now, error=str(exc))
        loop = asyncio.get_running_loop()
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
        sock.setblocking(False)
        try:
            sock.bind(("", 0))
            probe = _build_packet(msg_type=MSG_GET_SERVICE, payload=b"", tagged=True)
            start = loop.time()
            await loop.sock_sendto(sock, probe, (host, port))
            try:
                await asyncio.wait_for(loop.sock_recv(sock, 4096), timeout=1.5)
            except asyncio.TimeoutError:
                return DeviceHealth(
                    online=False,
                    last_checked=now,
                    error=f"No LIFX reply from {host}:{port} within 1.5s",
                )
            latency_ms = (loop.time() - start) * 1000.0
            return DeviceHealth(online=True, latency_ms=latency_ms, last_checked=now)
        except OSError as exc:
            return DeviceHealth(
                online=False,
                last_checked=now,
                error=f"LIFX probe failed for {host}: {exc}",
            )
        finally:
            sock.close()


# ============================================================================
# Broadcast discovery
# ============================================================================


async def discover_lifx_bulbs(timeout: float = 2.0) -> List[dict]:
    """Broadcast a GetService probe on the LAN and collect StateService replies.

    Returns ``[{"ip": ..., "mac": ..., "port": ...}, ...]``.
    """
    loop = asyncio.get_running_loop()
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    sock.setblocking(False)
    try:
        sock.bind(("", 0))
        probe = _build_packet(msg_type=MSG_GET_SERVICE, payload=b"", tagged=True)
        await loop.sock_sendto(sock, probe, ("255.255.255.255", LIFX_PORT))
        results: list[dict] = []
        seen_macs: set[str] = set()
        deadline = loop.time() + timeout
        while True:
            remaining = deadline - loop.time()
            if remaining <= 0:
                break
            try:
                raw, addr = await asyncio.wait_for(
                    loop.sock_recvfrom(sock, 4096),
                    timeout=remaining,
                )
            except asyncio.TimeoutError:
                break
            parsed = _parse_state_service_reply(raw)
            if not parsed:
                continue
            mac = parsed["mac"]
            if mac in seen_macs:
                continue
            seen_macs.add(mac)
            results.append(
                {
                    "ip": addr[0],
                    "mac": mac,
                    "port": parsed["port"],
                }
            )
        return results
    finally:
        sock.close()
