"""Adalight serial LED client — sends pixel data over serial using the Adalight protocol."""

import asyncio
import concurrent.futures
from datetime import datetime, timezone
from typing import Optional, Tuple

import numpy as np

from ledgrab.core.devices.led_client import DeviceHealth, LEDClient
from ledgrab.core.devices.serial_transport import (
    open_transport,
    parse_serial_url,
    port_exists,
)
from ledgrab.utils import get_logger

logger = get_logger(__name__)

ARDUINO_RESET_DELAY = 2.0  # seconds to wait after opening serial for Arduino bootloader


def parse_adalight_url(url: str) -> Tuple[str, int]:
    """Backwards-compatible alias for :func:`parse_serial_url`."""
    return parse_serial_url(url)


def _build_adalight_header(led_count: int) -> bytes:
    """Build the 6-byte Adalight protocol header.

    Format: 'A' 'd' 'a' <count_hi> <count_lo> <checksum>
    where count = led_count - 1 (zero-indexed).
    """
    count = led_count - 1
    hi = (count >> 8) & 0xFF
    lo = count & 0xFF
    checksum = hi ^ lo ^ 0x55
    return bytes([ord("A"), ord("d"), ord("a"), hi, lo, checksum])


class AdalightClient(LEDClient):
    """LED client for Arduino Adalight serial devices."""

    def __init__(self, url: str, led_count: int = 0, baud_rate: Optional[int] = None, **kwargs):
        """Initialize Adalight client.

        Args:
            url: Serial port string, e.g. "COM3" or "COM3:230400"
            led_count: Number of LEDs on the strip (required for Adalight header)
            baud_rate: Override baud rate (if None, parsed from url or default 115200)
        """
        self._port, url_baud = parse_adalight_url(url)
        self._baud_rate = baud_rate or url_baud
        self._led_count = led_count
        self._serial = None
        self._connected = False

        # Pre-compute Adalight header if led_count is known
        self._header = _build_adalight_header(led_count) if led_count > 0 else b""
        self._header_len = len(self._header)

        # Pre-allocated wire buffer (header + RGB payload).  Resized on the
        # first frame and reused thereafter so the hot path performs no
        # allocations — only a single memcpy of the pixel bytes.
        self._frame_buf: Optional[bytearray] = None
        self._frame_buf_n: int = 0
        # Scratch uint8 array used to coerce non-uint8 / non-contiguous input
        # without allocating a fresh array per frame.
        self._u8_scratch: Optional[np.ndarray] = None
        self._u8_scratch_n: int = 0
        # Dedicated single-worker executor for serial writes.  Using
        # ``loop.run_in_executor`` against this avoids the per-call
        # ``contextvars.copy_context()`` and ``functools.partial`` overhead
        # that ``asyncio.to_thread`` incurs (~5–10 µs per call), and
        # guarantees FIFO ordering of writes from this client even when
        # other tasks are using the default executor.
        self._tx_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None

    async def connect(self) -> bool:
        """Open serial port and wait for Arduino reset."""
        try:
            self._serial = open_transport(self._port, baud_rate=self._baud_rate, timeout=1)
            # Single-worker executor — created here so the thread is bound
            # to this client's lifecycle (started on connect, shut down on
            # close).  ``thread_name_prefix`` makes it identifiable in
            # diagnostics.
            self._tx_executor = concurrent.futures.ThreadPoolExecutor(
                max_workers=1,
                thread_name_prefix=f"adalight-tx-{self._port}",
            )
            await asyncio.get_running_loop().run_in_executor(self._tx_executor, self._serial.open)
            # Wait for Arduino to finish bootloader reset (non-blocking).
            # USB-to-TTL adapters without DTR don't reset, but the delay
            # is harmless on those — keeps the path uniform.
            await asyncio.sleep(ARDUINO_RESET_DELAY)
            self._connected = True
            logger.info(
                f"Adalight connected: {self._port} @ {self._baud_rate} baud "
                f"({self._led_count} LEDs)"
            )
            return True
        except Exception as e:
            logger.error(f"Failed to open serial port {self._port}: {e}")
            if self._tx_executor is not None:
                self._tx_executor.shutdown(wait=False)
                self._tx_executor = None
            raise RuntimeError(f"Failed to open serial port {self._port}: {e}")

    async def close(self) -> None:
        """Send black frame and close the serial port."""
        loop = asyncio.get_running_loop()
        executor = self._tx_executor
        if (
            self._connected
            and self._serial
            and self._serial.is_open
            and self._led_count > 0
            and executor is not None
        ):
            try:
                black = np.zeros((self._led_count, 3), dtype=np.uint8)
                frame = self._build_frame(black, brightness=255)
                logger.info(
                    f"Adalight sending black frame: {self._port} "
                    f"({self._led_count} LEDs, {len(frame)} bytes)"
                )
                await loop.run_in_executor(executor, self._serial.write, frame)
                await loop.run_in_executor(executor, self._serial.flush)
                logger.info(f"Adalight black frame sent and flushed: {self._port}")
            except Exception as e:
                logger.warning(f"Failed to send black frame on close: {e}")
        else:
            logger.warning(
                f"Adalight close skipped black frame: port={self._port} "
                f"connected={self._connected} serial={self._serial is not None} "
                f"is_open={self._serial.is_open if self._serial else 'N/A'} "
                f"led_count={self._led_count}"
            )
        self._connected = False
        if self._serial is not None:
            try:
                self._serial.close()
            except Exception as e:
                logger.warning(f"Error closing serial port: {e}")
        self._serial = None
        if self._tx_executor is not None:
            self._tx_executor.shutdown(wait=False)
            self._tx_executor = None
        logger.info(f"Adalight disconnected: {self._port}")

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

    async def send_pixels(
        self,
        pixels,
        brightness: int = 255,
    ) -> bool:
        """Send pixel data over serial using Adalight protocol (non-blocking).

        Args:
            pixels: numpy array (N, 3) uint8 or list of (R, G, B) tuples
            brightness: Global brightness (0-255)
        """
        executor = self._tx_executor
        if not self.is_connected or executor is None:
            return False

        try:
            frame = self._build_frame(pixels, brightness)
            # ``run_in_executor`` skips the per-call ``contextvars.copy_context``
            # / ``functools.partial`` overhead that ``asyncio.to_thread`` does.
            await asyncio.get_running_loop().run_in_executor(executor, self._serial.write, frame)
            return True
        except Exception as e:
            logger.error(f"Adalight send_pixels error: {e}")
            return False

    @property
    def supports_fast_send(self) -> bool:
        # Serial write is blocking — use async send_pixels path instead
        return False

    def _ensure_frame_buf(self, n_leds: int) -> None:
        """Lazily allocate / resize the wire-format frame buffer.

        Header bytes are written once at the front; subsequent calls only
        memcpy the pixel payload into the trailing slot.
        """
        needed = self._header_len + n_leds * 3
        if self._frame_buf is None or len(self._frame_buf) != needed:
            buf = bytearray(needed)
            buf[: self._header_len] = self._header
            self._frame_buf = buf
            self._frame_buf_n = n_leds

    def _ensure_u8_scratch(self, n_leds: int) -> np.ndarray:
        """Pre-allocated (N, 3) uint8 scratch for non-conforming inputs."""
        if self._u8_scratch is None or self._u8_scratch_n != n_leds:
            self._u8_scratch = np.empty((n_leds, 3), dtype=np.uint8)
            self._u8_scratch_n = n_leds
        return self._u8_scratch

    def _build_frame(self, pixels, brightness: int) -> bytes:
        """Build a complete Adalight frame in the pre-allocated wire buffer.

        The processor loop hands us a contiguous (N, 3) uint8 array with
        brightness already applied (see ``_cached_brightness``), so the hot
        path is one memcpy from the pixel buffer into the trailing slot of
        ``_frame_buf``.  All other input shapes (lists of tuples, wrong
        dtype, non-contiguous views) coerce into a pre-allocated uint8
        scratch before the memcpy — still allocation-free in steady state.
        """
        if isinstance(pixels, np.ndarray):
            n_leds = pixels.shape[0]
            if pixels.dtype == np.uint8 and pixels.flags["C_CONTIGUOUS"]:
                # Hot path: input matches wire format exactly.
                arr = pixels
            else:
                # Slow path: dtype mismatch or non-contiguous view.  Coerce
                # into a pre-allocated uint8 scratch.  Wider integer dtypes
                # are clamped to [0, 255] to match historical behaviour.
                arr = self._ensure_u8_scratch(n_leds)
                if pixels.dtype != np.uint8:
                    # Clamp wider integer dtypes to [0, 255] before the
                    # uint8 narrowing copy.  This is the rare slow path —
                    # one extra allocation here is fine.
                    np.copyto(arr, np.clip(pixels, 0, 255), casting="unsafe")
                else:
                    np.copyto(arr, pixels)
        else:
            # List/tuple input — rare path, only hit by tests/legacy callers.
            arr = np.array(pixels, dtype=np.uint8)
            n_leds = arr.shape[0]

        self._ensure_frame_buf(n_leds)
        # memcpy pixel bytes into the trailing slot of the pre-built buffer
        view = memoryview(self._frame_buf)
        view[self._header_len :] = memoryview(arr).cast("B")
        return self._frame_buf

    @classmethod
    async def check_health(
        cls,
        url: str,
        http_client,
        prev_health: Optional[DeviceHealth] = None,
    ) -> DeviceHealth:
        """Check if the serial port exists without opening it.

        Enumerates COM ports (or USB devices on Android) to avoid
        exclusive-access conflicts on Windows.
        """
        port, _baud = parse_adalight_url(url)

        try:
            if port_exists(port):
                return DeviceHealth(
                    online=True,
                    latency_ms=0.0,
                    last_checked=datetime.now(timezone.utc),
                    device_name=prev_health.device_name if prev_health else None,
                    device_version=None,
                    device_led_count=prev_health.device_led_count if prev_health else None,
                )
            return DeviceHealth(
                online=False,
                last_checked=datetime.now(timezone.utc),
                error=f"Serial port {port} not found",
            )
        except Exception as e:
            return DeviceHealth(
                online=False,
                last_checked=datetime.now(timezone.utc),
                error=str(e),
            )
