"""Screen capture functionality using mss library."""

from dataclasses import dataclass
from typing import List

import mss
import numpy as np

from wled_controller.utils import get_logger, get_monitor_names, get_monitor_refresh_rates

logger = get_logger(__name__)


@dataclass
class DisplayInfo:
    """Information about a display/monitor."""

    index: int
    name: str
    width: int
    height: int
    x: int
    y: int
    is_primary: bool
    refresh_rate: int  # in Hz


@dataclass
class ScreenCapture:
    """Captured screen image data."""

    image: np.ndarray
    width: int
    height: int
    display_index: int


@dataclass
class BorderPixels:
    """Border pixels extracted from screen edges."""

    top: np.ndarray
    right: np.ndarray
    bottom: np.ndarray
    left: np.ndarray


def get_available_displays() -> List[DisplayInfo]:
    """Get list of available displays/monitors.

    Returns:
        List of DisplayInfo objects for each available monitor

    Raises:
        RuntimeError: If unable to detect displays
    """
    try:
        # Get friendly monitor names (Windows only, falls back to generic names)
        monitor_names = get_monitor_names()

        # Get monitor refresh rates (Windows only, falls back to 60Hz)
        refresh_rates = get_monitor_refresh_rates()

        with mss.mss() as sct:
            displays = []

            # Skip the first monitor (combined virtual screen on multi-monitor setups)
            for idx, monitor in enumerate(sct.monitors[1:], start=0):
                # Use friendly name from WMI if available, otherwise generic name
                friendly_name = monitor_names.get(idx, f"Display {idx}")

                # Use detected refresh rate or default to 60Hz
                refresh_rate = refresh_rates.get(idx, 60)

                display_info = 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,
                )
                displays.append(display_info)

            logger.info(f"Detected {len(displays)} display(s)")
            return displays

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


def capture_display(display_index: int = 0) -> ScreenCapture:
    """Capture the specified display.

    Args:
        display_index: Index of the display to capture (0-based)

    Returns:
        ScreenCapture object containing the captured image

    Raises:
        ValueError: If display_index is invalid
        RuntimeError: If screen capture fails
    """
    try:
        with mss.mss() as sct:
            # mss monitors[0] is the combined screen, monitors[1+] are individual displays
            monitor_index = display_index + 1

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

            monitor = sct.monitors[monitor_index]

            # Capture screenshot
            screenshot = 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"Captured display {display_index}: {monitor['width']}x{monitor['height']}"
            )

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

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


def extract_border_pixels(
    screen_capture: ScreenCapture,
    border_width: int = 10
) -> BorderPixels:
    """Extract border pixels from screen capture.

    Args:
        screen_capture: Captured screen image
        border_width: Width of the border in pixels to extract

    Returns:
        BorderPixels object containing pixels from each edge

    Raises:
        ValueError: If border_width is invalid
    """
    if border_width < 1:
        raise ValueError("border_width must be at least 1")

    if border_width > min(screen_capture.width, screen_capture.height) // 4:
        raise ValueError(
            f"border_width {border_width} is too large for screen size "
            f"{screen_capture.width}x{screen_capture.height}"
        )

    img = screen_capture.image
    height, width = img.shape[:2]

    # Extract border regions
    # Top edge: top border_width rows, full width
    top = img[:border_width, :, :]

    # Bottom edge: bottom border_width rows, full width
    bottom = img[-border_width:, :, :]

    # Right edge: right border_width columns, full height
    right = img[:, -border_width:, :]

    # Left edge: left border_width columns, full height
    left = img[:, :border_width, :]

    logger.debug(
        f"Extracted borders: top={top.shape}, right={right.shape}, "
        f"bottom={bottom.shape}, left={left.shape}"
    )

    return BorderPixels(
        top=top,
        right=right,
        bottom=bottom,
        left=left,
    )


def get_edge_segments(
    edge_pixels: np.ndarray,
    segment_count: int,
    edge_name: str
) -> List[np.ndarray]:
    """Divide edge pixels into segments.

    Args:
        edge_pixels: Pixel array for one edge
        segment_count: Number of segments to divide into
        edge_name: Name of the edge (for orientation)

    Returns:
        List of pixel arrays, one per segment

    Raises:
        ValueError: If segment_count is invalid
    """
    if segment_count < 1:
        raise ValueError("segment_count must be at least 1")

    # Determine the dimension to divide
    # For top/bottom edges: divide along width (axis 1)
    # For left/right edges: divide along height (axis 0)
    if edge_name in ["top", "bottom"]:
        divide_axis = 1  # Width
        edge_length = edge_pixels.shape[1]
    else:  # left, right
        divide_axis = 0  # Height
        edge_length = edge_pixels.shape[0]

    # Use float stepping so multiple LEDs can share pixels when
    # segment_count > edge_length (e.g. after downscaling).
    step = edge_length / segment_count

    segments = []
    for i in range(segment_count):
        start = int(i * step)
        end = max(start + 1, int((i + 1) * step))
        # Clamp to edge bounds
        end = min(end, edge_length)

        if divide_axis == 1:
            segment = edge_pixels[:, start:end, :]
        else:
            segment = edge_pixels[start:end, :, :]

        segments.append(segment)

    return segments


def calculate_average_color(pixels: np.ndarray) -> tuple[int, int, int]:
    """Calculate average color of a pixel region.

    Args:
        pixels: Pixel array (height, width, 3)

    Returns:
        Tuple of (R, G, B) average values
    """
    if pixels.size == 0:
        return (0, 0, 0)

    # Calculate mean across height and width dimensions
    mean_color = np.mean(pixels, axis=(0, 1))

    # Convert to integers and clamp to valid range
    r = int(np.clip(mean_color[0], 0, 255))
    g = int(np.clip(mean_color[1], 0, 255))
    b = int(np.clip(mean_color[2], 0, 255))

    return (r, g, b)


def calculate_median_color(pixels: np.ndarray) -> tuple[int, int, int]:
    """Calculate median color of a pixel region.

    Args:
        pixels: Pixel array (height, width, 3)

    Returns:
        Tuple of (R, G, B) median values
    """
    if pixels.size == 0:
        return (0, 0, 0)

    # Calculate median across height and width dimensions
    median_color = np.median(pixels, axis=(0, 1))

    # Convert to integers and clamp to valid range
    r = int(np.clip(median_color[0], 0, 255))
    g = int(np.clip(median_color[1], 0, 255))
    b = int(np.clip(median_color[2], 0, 255))

    return (r, g, b)


def calculate_dominant_color(pixels: np.ndarray) -> tuple[int, int, int]:
    """Calculate dominant color of a pixel region using simple clustering.

    Args:
        pixels: Pixel array (height, width, 3)

    Returns:
        Tuple of (R, G, B) dominant color values
    """
    if pixels.size == 0:
        return (0, 0, 0)

    # Reshape to (n_pixels, 3)
    pixels_reshaped = pixels.reshape(-1, 3)

    # For performance, sample pixels if there are too many
    max_samples = 1000
    if len(pixels_reshaped) > max_samples:
        indices = np.random.choice(len(pixels_reshaped), max_samples, replace=False)
        pixels_reshaped = pixels_reshaped[indices]

    # Simple dominant color: quantize colors and find most common
    # Reduce color space to 32 levels per channel for binning
    quantized = (pixels_reshaped // 8) * 8

    # Find unique colors and their counts
    unique_colors, counts = np.unique(quantized, axis=0, return_counts=True)

    # Get the most common color
    dominant_idx = np.argmax(counts)
    dominant_color = unique_colors[dominant_idx]

    r = int(np.clip(dominant_color[0], 0, 255))
    g = int(np.clip(dominant_color[1], 0, 255))
    b = int(np.clip(dominant_color[2], 0, 255))

    return (r, g, b)
