"""MSS-based screen capture engine (cross-platform)."""

from typing import Any, Dict, List, Optional

import mss
import numpy as np

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

logger = get_logger(__name__)


class MSSCaptureStream(CaptureStream):
    """MSS capture stream for a specific display."""

    def __init__(self, display_index: int, config: Dict[str, Any]):
        super().__init__(display_index, config)
        self._sct = None

    def initialize(self) -> None:
        try:
            self._sct = mss.mss()
            self._initialized = True
            logger.info(f"MSS capture stream initialized (display={self.display_index})")
        except Exception as e:
            raise RuntimeError(f"Failed to initialize MSS: {e}")

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

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

        try:
            # mss monitors[0] is the combined screen, monitors[1+] are individual displays
            monitor_index = self.display_index + 1

            if monitor_index >= len(self._sct.monitors):
                raise ValueError(
                    f"Invalid display index {self.display_index}. "
                    f"Available displays: 0-{len(self._sct.monitors) - 2}"
                )

            monitor = self._sct.monitors[monitor_index]
            screenshot = self._sct.grab(monitor)

            # Direct bytes→numpy (skips PIL intermediate object)
            img_array = np.frombuffer(
                screenshot.rgb, dtype=np.uint8,
            ).reshape(screenshot.height, screenshot.width, 3)

            logger.debug(
                f"MSS captured display {self.display_index}: {monitor['width']}x{monitor['height']}"
            )

            return ScreenCapture(
                image=img_array,
                width=monitor["width"],
                height=monitor["height"],
                display_index=self.display_index,
            )

        except ValueError:
            raise
        except Exception as e:
            logger.error(f"Failed to capture display {self.display_index} with MSS: {e}")
            raise RuntimeError(f"Screen capture failed: {e}")


class MSSEngine(CaptureEngine):
    """MSS-based screen capture engine.

    Uses the mss library for cross-platform screen capture support.
    Works on Windows, macOS, and Linux.
    """

    ENGINE_TYPE = "mss"
    ENGINE_PRIORITY = 1

    @classmethod
    def is_available(cls) -> bool:
        try:
            import importlib.util
            return importlib.util.find_spec("mss") is not None
        except (ImportError, ModuleNotFoundError, ValueError):
            return False

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

    @classmethod
    def get_available_displays(cls) -> List[DisplayInfo]:
        try:
            monitor_names = get_monitor_names()
            refresh_rates = get_monitor_refresh_rates()

            displays = []
            with mss.mss() as sct:
                for idx, monitor in enumerate(sct.monitors[1:], start=0):
                    friendly_name = monitor_names.get(idx, f"Display {idx}")
                    refresh_rate = refresh_rates.get(idx, 60)

                    displays.append(DisplayInfo(
                        index=idx,
                        name=friendly_name,
                        width=monitor["width"],
                        height=monitor["height"],
                        x=monitor["left"],
                        y=monitor["top"],
                        is_primary=(idx == 0),
                        refresh_rate=refresh_rate,
                    ))

            logger.debug(f"MSS detected {len(displays)} display(s)")
            return displays

        except Exception as e:
            logger.error(f"Failed to detect displays with MSS: {e}")
            raise RuntimeError(f"Failed to detect displays: {e}")

    @classmethod
    def create_stream(cls, display_index: int, config: Dict[str, Any]) -> MSSCaptureStream:
        return MSSCaptureStream(display_index, config)
