"""Group LED client — distributes pixels to multiple child devices."""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np

from ledgrab.core.devices.led_client import LEDClient, ProviderDeps, create_led_client
from ledgrab.utils import get_logger

if TYPE_CHECKING:
    from ledgrab.core.devices.device_config import GroupConfig

logger = get_logger(__name__)


class GroupLEDClient(LEDClient):
    """Virtual LED client that fans out pixels to multiple child devices.

    Sequence mode: slices the pixel array by each child's LED count.
    Independent mode: resamples the full pixel array to each child's LED count.
    """

    def __init__(self, config: "GroupConfig", deps: "ProviderDeps"):
        self._device_store = deps.device_store
        self._deps = deps
        self._device_id = config.device_id
        self._group_mode = config.group_mode
        self._group_device_ids = list(config.group_device_ids)
        # Populated on connect()
        self._children: List[Tuple[LEDClient, int]] = []  # (client, led_count)
        self._connected = False
        self._total_led_count = 0

    async def connect(self) -> bool:
        """Recursively create and connect child clients."""
        connected_clients: List[LEDClient] = []
        try:
            for child_id in self._group_device_ids:
                device = self._device_store.get_device(child_id)
                child_config = device.to_config()
                client = create_led_client(child_config, deps=self._deps)
                await client.connect()
                connected_clients.append(client)

                # Use client-reported LED count if available
                client_leds = client.device_led_count
                led_count = client_leds if client_leds and client_leds > 0 else device.led_count
                self._children.append((client, led_count))

            self._total_led_count = sum(lc for _, lc in self._children)
            self._connected = True
            logger.info(
                f"Group {self._device_id} connected: {len(self._children)} children, "
                f"{self._total_led_count} total LEDs ({self._group_mode} mode)"
            )
            return True
        except Exception:
            # Clean up already-connected clients on failure
            for c in connected_clients:
                try:
                    await c.close()
                except Exception:
                    pass
            self._children.clear()
            raise

    async def close(self) -> None:
        """Close all child clients."""
        for client, _ in self._children:
            try:
                await client.close()
            except Exception as e:
                logger.warning(f"Error closing child client in group {self._device_id}: {e}")
        self._children.clear()
        self._connected = False

    @property
    def is_connected(self) -> bool:
        return self._connected and all(c.is_connected for c, _ in self._children)

    @property
    def device_led_count(self) -> Optional[int]:
        if self._group_mode == "sequence":
            return self._total_led_count
        return None  # independent mode uses user-specified led_count

    async def send_pixels(
        self,
        pixels: Union[List[Tuple[int, int, int]], np.ndarray],
        brightness: int = 255,
    ) -> bool:
        if not self._children:
            return False

        pixel_array = np.asarray(pixels, dtype=np.uint8)

        if self._group_mode == "sequence":
            return await self._send_sequence(pixel_array, brightness)
        else:
            return await self._send_independent(pixel_array, brightness)

    async def _send_sequence(self, pixels: np.ndarray, brightness: int) -> bool:
        """Slice pixel array and distribute sequentially to children."""
        tasks = []
        offset = 0
        for client, led_count in self._children:
            end = offset + led_count
            chunk = (
                pixels[offset:end]
                if offset < len(pixels)
                else np.zeros((led_count, 3), dtype=np.uint8)
            )
            # Pad with black if chunk is shorter than expected
            if len(chunk) < led_count:
                chunk = np.vstack(
                    [
                        chunk,
                        np.zeros((led_count - len(chunk), 3), dtype=np.uint8),
                    ]
                )
            tasks.append(client.send_pixels(chunk, brightness))
            offset = end
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return all(r is True for r in results if not isinstance(r, Exception))

    async def _send_independent(self, pixels: np.ndarray, brightness: int) -> bool:
        """Resample full pixel array independently to each child's LED count."""
        n_src = len(pixels)
        tasks = []
        if n_src < 2:
            # Single pixel or empty — broadcast same color to all
            color = pixels[0] if n_src == 1 else np.array([0, 0, 0], dtype=np.uint8)
            for client, led_count in self._children:
                tiled = np.tile(color, (led_count, 1))
                tasks.append(client.send_pixels(tiled, brightness))
        else:
            src_indices = np.linspace(0, 1, n_src)
            for client, led_count in self._children:
                dst_indices = np.linspace(0, 1, led_count)
                resampled = np.column_stack(
                    [
                        np.interp(dst_indices, src_indices, pixels[:, ch])
                        for ch in range(pixels.shape[1])
                    ]
                ).astype(np.uint8)
                tasks.append(client.send_pixels(resampled, brightness))
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return all(r is True for r in results if not isinstance(r, Exception))

    async def snapshot_device_state(self) -> Optional[dict]:
        """Snapshot all children's states."""
        states = {}
        for i, (client, _) in enumerate(self._children):
            state = await client.snapshot_device_state()
            if state is not None:
                states[i] = state
        return states if states else None

    async def restore_device_state(self, state: Optional[dict]) -> None:
        """Restore all children's states."""
        if not state:
            return
        for i, (client, _) in enumerate(self._children):
            child_state = state.get(i)
            if child_state is not None:
                await client.restore_device_state(child_state)
