"""Demo capture engine — virtual displays with animated test patterns."""

import time
from typing import Any, Dict, List, Optional

import numpy as np

from wled_controller.config import is_demo_mode
from wled_controller.core.capture_engines.base import (
    CaptureEngine,
    CaptureStream,
    DisplayInfo,
    ScreenCapture,
)
from wled_controller.utils import get_logger

logger = get_logger(__name__)

# Virtual display definitions: (name, width, height, x, y, is_primary)
_VIRTUAL_DISPLAYS = [
    ("Demo Display 1080p", 1920, 1080, 0, 360, True),
    ("Demo Ultrawide", 3440, 1440, 1920, 0, False),
    ("Demo Portrait", 1080, 1920, 5360, 0, False),
]


class DemoCaptureStream(CaptureStream):
    """Demo capture stream producing a radial rainbow centred on the screen.

    The rainbow rotates slowly over time — hue is mapped to the angle from
    the screen centre, and brightness fades toward the edges.
    """

    _RENDER_SCALE = 4  # render at 1/4 resolution, then upscale

    def __init__(self, display_index: int, config: Dict[str, Any]):
        super().__init__(display_index, config)
        self._width: int = config.get("width", 1920)
        self._height: int = config.get("height", 1080)
        # Pre-compute at render resolution
        rw = max(1, self._width // self._RENDER_SCALE)
        rh = max(1, self._height // self._RENDER_SCALE)
        self._rw = rw
        self._rh = rh
        # Coordinate grids centred at (0, 0), aspect-corrected so the
        # gradient is circular even on non-square displays
        aspect = self._width / max(self._height, 1)
        x = np.linspace(-aspect, aspect, rw, dtype=np.float32)
        y = np.linspace(-1.0, 1.0, rh, dtype=np.float32)
        self._yy, self._xx = np.meshgrid(y, x, indexing="ij")
        # Pre-compute angle (atan2) and radius — they don't change per frame
        self._angle = np.arctan2(self._yy, self._xx)  # -pi..pi
        self._radius = np.sqrt(self._xx ** 2 + self._yy ** 2)

    def initialize(self) -> None:
        self._initialized = True
        logger.info(
            f"Demo capture stream initialized "
            f"(display={self.display_index}, {self._width}x{self._height})"
        )

    def cleanup(self) -> None:
        self._initialized = False
        logger.info(f"Demo capture stream cleaned up (display={self.display_index})")

    def capture_frame(self) -> Optional[ScreenCapture]:
        if not self._initialized:
            self.initialize()

        t = time.time() % 1e6

        # Hue = angle from centre, rotating over time
        rotation = t * 0.15  # radians per second
        hue = ((self._angle + rotation) / (2.0 * np.pi)) % 1.0

        # Saturation: full

        # Value: bright at centre, fading toward edges
        max_r = float(self._radius.max()) or 1.0
        val = np.clip(1.0 - 0.6 * (self._radius / max_r), 0.0, 1.0)

        # Vectorised HSV → RGB (S=1 simplification)
        h6 = hue * 6.0
        sector = h6.astype(np.int32) % 6
        frac = h6 - np.floor(h6)
        q = val * (1.0 - frac)
        t_ch = val * frac  # "t" channel in HSV conversion

        r = np.where(sector == 0, val,
            np.where(sector == 1, q,
            np.where(sector == 2, 0,
            np.where(sector == 3, 0,
            np.where(sector == 4, t_ch, val)))))
        g = np.where(sector == 0, t_ch,
            np.where(sector == 1, val,
            np.where(sector == 2, val,
            np.where(sector == 3, q,
            np.where(sector == 4, 0, 0)))))
        b = np.where(sector == 0, 0,
            np.where(sector == 1, 0,
            np.where(sector == 2, t_ch,
            np.where(sector == 3, val,
            np.where(sector == 4, val, q)))))

        small_u8 = (np.stack([r, g, b], axis=-1) * 255.0).astype(np.uint8)

        # Upscale to full resolution
        if self._RENDER_SCALE > 1:
            image = np.repeat(
                np.repeat(small_u8, self._RENDER_SCALE, axis=0),
                self._RENDER_SCALE, axis=1,
            )[: self._height, : self._width]
        else:
            image = small_u8

        return ScreenCapture(
            image=image,
            width=self._width,
            height=self._height,
            display_index=self.display_index,
        )


class DemoCaptureEngine(CaptureEngine):
    """Virtual capture engine for demo mode.

    Provides virtual displays and produces animated test-pattern frames
    so the full capture pipeline works without real monitors.
    """

    ENGINE_TYPE = "demo"
    ENGINE_PRIORITY = 1000  # Highest priority in demo mode

    @classmethod
    def is_available(cls) -> bool:
        return is_demo_mode()

    @classmethod
    def get_default_config(cls) -> Dict[str, Any]:
        return {}

    @classmethod
    def get_available_displays(cls) -> List[DisplayInfo]:
        displays = []
        for idx, (name, width, height, x, y, primary) in enumerate(_VIRTUAL_DISPLAYS):
            displays.append(DisplayInfo(
                index=idx,
                name=name,
                width=width,
                height=height,
                x=x,
                y=y,
                is_primary=primary,
                refresh_rate=60,
            ))
        logger.debug(f"Demo engine: {len(displays)} virtual display(s)")
        return displays

    @classmethod
    def create_stream(
        cls, display_index: int, config: Dict[str, Any],
    ) -> DemoCaptureStream:
        if display_index < 0 or display_index >= len(_VIRTUAL_DISPLAYS):
            raise ValueError(
                f"Invalid demo display index {display_index}. "
                f"Available: 0-{len(_VIRTUAL_DISPLAYS) - 1}"
            )
        name, width, height, *_ = _VIRTUAL_DISPLAYS[display_index]
        stream_config = {**config, "width": width, "height": height}
        return DemoCaptureStream(display_index, stream_config)
