"""Govee LAN API LED client.

Govee opened a local LAN API in 2023 for its Wi-Fi smart bulbs and
ambient-light kits. Discovery is multicast on ``239.255.255.250:4001``;
control commands go unicast to the bulb's port ``4003``; the bulb sends
responses on port ``4002`` (which we don't listen on for ambient streaming).

Prerequisite for every device: the user must toggle "LAN Control" ON in the
Govee Home app (Device → ⚙ → LAN Control). Devices with LAN Control disabled
do not respond to discovery or commands. The UI hint copy reminds the user.

URL scheme: ``govee://<host>`` or bare ``<host>``. Port 4003 is fixed.

Reference: https://app-h5.govee.com/user-manual/wlan-guide
"""

from __future__ import annotations

import asyncio
import json
import socket
import time
from datetime import datetime, timezone
from typing import List, Optional, Tuple, Union
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__)

GOVEE_DISCOVERY_PORT = 4001
GOVEE_RESPONSE_PORT = 4002
GOVEE_CONTROL_PORT = 4003
GOVEE_MULTICAST_GROUP = "239.255.255.250"
DEFAULT_MIN_INTERVAL_S = 0.05  # ~20 Hz; UDP fire-and-forget, no ack

_DISCOVERY_REQUEST = json.dumps(
    {"msg": {"cmd": "scan", "data": {"account_topic": "reserve"}}}
).encode("utf-8")


# Govee mandates listening on port 4002 for scan replies; the OS only lets
# one process / socket own it at a time. Serialize concurrent discover
# calls (e.g. two browser tabs hitting /discover at once) so the second
# caller waits its turn instead of getting an empty result.
_DISCOVERY_LOCK: asyncio.Lock | None = None


def _discovery_lock() -> asyncio.Lock:
    """Return a module-level asyncio.Lock, created lazily.

    Can't construct at import-time because asyncio.Lock binds to the
    running event loop and module import may precede loop creation.
    """
    global _DISCOVERY_LOCK
    if _DISCOVERY_LOCK is None:
        _DISCOVERY_LOCK = asyncio.Lock()
    return _DISCOVERY_LOCK


def parse_govee_url(url: str) -> str:
    """Pull the host out of ``govee://host`` or bare ``host``."""
    if not url:
        raise ValueError("Govee URL is empty")
    raw = url.strip()
    if "://" in raw:
        parsed = urlparse(raw)
        host = parsed.hostname or ""
    else:
        parsed = urlparse(f"govee://{raw}")
        host = parsed.hostname or ""
    if not host:
        raise ValueError(f"Govee URL has no host: {url!r}")
    return host


class _GoveeProtocol(asyncio.DatagramProtocol):
    """Write-only datagram protocol. Bulb replies (on 4002) are not collected here."""

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

    def datagram_received(self, data, addr):
        pass

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


class GoveeClient(LEDClient):
    """LEDClient for a single Govee LAN-enabled bulb / ambient kit."""

    def __init__(
        self,
        url: str,
        led_count: int = 1,
        *,
        min_interval_s: float = DEFAULT_MIN_INTERVAL_S,
    ):
        self._host = parse_govee_url(url)
        self._port = GOVEE_CONTROL_PORT
        self._led_count = led_count
        self._min_interval_s = max(0.0, min_interval_s)
        self._transport: Optional[asyncio.DatagramTransport] = None
        self._protocol: Optional[_GoveeProtocol] = None
        self._connected = False
        self._next_tx_at: float = 0.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) -> Optional[int]:
        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(
                _GoveeProtocol, remote_addr=(self._host, self._port)
            )
        except OSError as exc:
            raise RuntimeError(f"Failed to open UDP to Govee at {self._host}: {exc}") from exc
        self._transport = transport
        self._protocol = protocol  # type: ignore[assignment]
        self._connected = True
        logger.info("GoveeClient 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 _send_json(self, payload: dict) -> None:
        assert self._transport is not None
        raw = json.dumps(payload, separators=(",", ":")).encode("utf-8")
        self._transport.sendto(raw)

    @staticmethod
    def _build_color_payload(r: int, g: int, b: int) -> dict:
        """Govee colorwc command. ``colorTemInKelvin=0`` selects pure RGB."""
        return {
            "msg": {
                "cmd": "colorwc",
                "data": {
                    "color": {"r": r & 0xFF, "g": g & 0xFF, "b": b & 0xFF},
                    "colorTemInKelvin": 0,
                },
            }
        }

    @staticmethod
    def _build_brightness_payload(value_0_100: int) -> dict:
        return {
            "msg": {
                "cmd": "brightness",
                "data": {"value": max(1, min(100, value_0_100))},
            }
        }

    @staticmethod
    def _build_power_payload(on: bool) -> dict:
        return {"msg": {"cmd": "turn", "data": {"value": 1 if on else 0}}}

    async def send_pixels(
        self,
        pixels: Union[List[Tuple[int, int, int]], np.ndarray],
        brightness: int = 255,
    ) -> bool:
        """Average the strip → colorwc with the resulting RGB."""
        if not self.is_connected:
            raise RuntimeError("GoveeClient 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)
        self._send_json(self._build_color_payload(r, g, b))
        self._next_tx_at = now + self._min_interval_s
        return True

    def send_pixels_fast(
        self,
        pixels: Union[List[Tuple[int, int, int]], np.ndarray],
        brightness: int = 255,
    ) -> None:
        """Synchronous variant for the hot loop."""
        if not self.is_connected or self._transport is None:
            raise RuntimeError("GoveeClient 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)
        self._send_json(self._build_color_payload(r, g, b))
        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("GoveeClient not connected")
        self._send_json(self._build_color_payload(r, g, b))

    async def set_brightness(self, brightness_0_100: int) -> None:
        if not self.is_connected:
            raise RuntimeError("GoveeClient not connected")
        self._send_json(self._build_brightness_payload(brightness_0_100))

    async def set_power(self, on: bool) -> None:
        if not self.is_connected:
            raise RuntimeError("GoveeClient not connected")
        self._send_json(self._build_power_payload(on))

    @classmethod
    async def check_health(
        cls,
        url: str,
        http_client,
        prev_health: Optional[DeviceHealth] = None,
    ) -> DeviceHealth:
        """Send devStatus and wait briefly for a reply on port 4002.

        Govee bulbs send responses to whatever port the request came from
        when using ``connected`` UDP, so we bind to a random ephemeral port
        and accept any reply. Health is best-effort — a silent bulb may
        still be online (it just hasn't toggled LAN Control on yet).
        """
        now = datetime.now(timezone.utc)
        try:
            host = parse_govee_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 = json.dumps({"msg": {"cmd": "devStatus", "data": {}}}).encode("utf-8")
            start = loop.time()
            await loop.sock_sendto(sock, probe, (host, GOVEE_CONTROL_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 Govee reply from {host} within 1.5s — is "
                        "LAN Control enabled in the Govee Home app?"
                    ),
                )
            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"Govee probe failed for {host}: {exc}",
            )
        finally:
            sock.close()


# ============================================================================
# Multicast discovery
# ============================================================================


def _parse_scan_reply(raw: bytes) -> Optional[dict]:
    """Parse a Govee scan reply into a flat metadata dict.

    Govee sends ``{"msg": {"cmd": "scan", "data": {"ip": ..., "device": ...,
    "sku": ..., "wifiVersionSoft": ..., ...}}}``. Returns the inner ``data``
    dict, or ``None`` for malformed packets.
    """
    try:
        payload = json.loads(raw.decode("utf-8", errors="replace"))
    except (json.JSONDecodeError, UnicodeDecodeError):
        return None
    if not isinstance(payload, dict):
        return None
    msg = payload.get("msg")
    if not isinstance(msg, dict):
        return None
    if msg.get("cmd") != "scan":
        return None
    data = msg.get("data")
    if not isinstance(data, dict):
        return None
    return data


async def discover_govee_devices(timeout: float = 2.0) -> List[dict]:
    """Multicast a scan request and collect Govee scan replies.

    Returns a list of ``{"ip": ..., "device": ..., "sku": ..., "version": ...}``
    dicts. Multicast / receive failures (no network, firewall, no LAN-enabled
    bulbs) yield an empty list rather than raising.

    Concurrent invocations are serialized via a module-level
    ``asyncio.Lock`` so two callers don't race on the protocol-mandated
    response port 4002.
    """
    async with _discovery_lock():
        return await _discover_govee_devices_locked(timeout)


async def _discover_govee_devices_locked(timeout: float) -> List[dict]:
    loop = asyncio.get_running_loop()
    # We bind a separate socket to the response port (4002). The
    # _DISCOVERY_LOCK serializes our own scans; the bind can still fail
    # if a non-LedGrab Govee tool on the same host owns 4002.
    recv_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    recv_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    recv_sock.setblocking(False)
    send_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    send_sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 2)
    send_sock.setblocking(False)
    try:
        try:
            recv_sock.bind(("", GOVEE_RESPONSE_PORT))
        except OSError as exc:
            logger.warning(
                "Govee discovery: cannot bind port %d (%s). Another Govee "
                "tool on this host may own the port; close it and retry.",
                GOVEE_RESPONSE_PORT,
                exc,
            )
            return []
        send_sock.bind(("", 0))
        await loop.sock_sendto(
            send_sock, _DISCOVERY_REQUEST, (GOVEE_MULTICAST_GROUP, GOVEE_DISCOVERY_PORT)
        )
        results: list[dict] = []
        seen_ips: 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(recv_sock, 4096),
                    timeout=remaining,
                )
            except asyncio.TimeoutError:
                break
            data = _parse_scan_reply(raw)
            if not data:
                continue
            ip = data.get("ip") or addr[0]
            if not ip or ip in seen_ips:
                continue
            seen_ips.add(ip)
            results.append(
                {
                    "ip": ip,
                    "device": data.get("device", ""),
                    "sku": data.get("sku", ""),
                    "version": data.get("wifiVersionSoft", "") or data.get("bleVersionSoft", ""),
                }
            )
        return results
    finally:
        recv_sock.close()
        send_sock.close()
