"""Calibration system for mapping screen pixels to LED positions."""

from dataclasses import dataclass, field
from typing import Dict, List, Literal, Tuple

from wled_controller.core.screen_capture import (
    BorderPixels,
    get_edge_segments,
    calculate_average_color,
    calculate_median_color,
    calculate_dominant_color,
)
from wled_controller.utils import get_logger

logger = get_logger(__name__)

# Edge traversal order for each (start_position, layout) combination.
# Determines which edge comes first when walking around the screen.
EDGE_ORDER: Dict[Tuple[str, str], List[str]] = {
    ("bottom_left", "clockwise"): ["left", "top", "right", "bottom"],
    ("bottom_left", "counterclockwise"): ["bottom", "right", "top", "left"],
    ("bottom_right", "clockwise"): ["bottom", "left", "top", "right"],
    ("bottom_right", "counterclockwise"): ["right", "top", "left", "bottom"],
    ("top_left", "clockwise"): ["top", "right", "bottom", "left"],
    ("top_left", "counterclockwise"): ["left", "bottom", "right", "top"],
    ("top_right", "clockwise"): ["right", "bottom", "left", "top"],
    ("top_right", "counterclockwise"): ["top", "left", "bottom", "right"],
}

# Whether LEDs are reversed on each edge for each (start_position, layout) combination.
EDGE_REVERSE: Dict[Tuple[str, str], Dict[str, bool]] = {
    ("bottom_left", "clockwise"): {"left": True, "top": False, "right": False, "bottom": True},
    ("bottom_left", "counterclockwise"): {"bottom": False, "right": True, "top": True, "left": False},
    ("bottom_right", "clockwise"): {"bottom": True, "left": True, "top": False, "right": False},
    ("bottom_right", "counterclockwise"): {"right": True, "top": True, "left": False, "bottom": False},
    ("top_left", "clockwise"): {"top": False, "right": False, "bottom": True, "left": True},
    ("top_left", "counterclockwise"): {"left": False, "bottom": False, "right": True, "top": True},
    ("top_right", "clockwise"): {"right": False, "bottom": True, "left": True, "top": False},
    ("top_right", "counterclockwise"): {"top": True, "left": False, "bottom": False, "right": True},
}


@dataclass
class CalibrationSegment:
    """Configuration for one segment of the LED strip."""

    edge: Literal["top", "right", "bottom", "left"]
    led_start: int
    led_count: int
    reverse: bool = False


@dataclass
class CalibrationConfig:
    """Complete calibration configuration.

    Stores only the core parameters. Segments (with led_start, reverse, edge order)
    are derived at runtime via the `segments` property.
    """

    layout: Literal["clockwise", "counterclockwise"]
    start_position: Literal["top_left", "top_right", "bottom_left", "bottom_right"]
    offset: int = 0
    leds_top: int = 0
    leds_right: int = 0
    leds_bottom: int = 0
    leds_left: int = 0
    # Per-edge span: fraction of screen side covered by LEDs (0.0–1.0)
    span_top_start: float = 0.0
    span_top_end: float = 1.0
    span_right_start: float = 0.0
    span_right_end: float = 1.0
    span_bottom_start: float = 0.0
    span_bottom_end: float = 1.0
    span_left_start: float = 0.0
    span_left_end: float = 1.0

    def build_segments(self) -> List[CalibrationSegment]:
        """Derive segment list from core parameters."""
        key = (self.start_position, self.layout)
        edge_order = EDGE_ORDER.get(key, ["bottom", "right", "top", "left"])
        reverse_map = EDGE_REVERSE.get(key, {})

        led_counts = {
            "top": self.leds_top,
            "right": self.leds_right,
            "bottom": self.leds_bottom,
            "left": self.leds_left,
        }

        segments = []
        led_start = 0
        for edge in edge_order:
            count = led_counts[edge]
            if count > 0:
                segments.append(CalibrationSegment(
                    edge=edge,
                    led_start=led_start,
                    led_count=count,
                    reverse=reverse_map.get(edge, False),
                ))
                led_start += count

        return segments

    @property
    def segments(self) -> List[CalibrationSegment]:
        """Get derived segment list."""
        return self.build_segments()

    def get_edge_span(self, edge: str) -> tuple[float, float]:
        """Get span (start, end) for a given edge."""
        return (
            getattr(self, f"span_{edge}_start", 0.0),
            getattr(self, f"span_{edge}_end", 1.0),
        )

    def validate(self) -> bool:
        """Validate calibration configuration.

        Returns:
            True if configuration is valid

        Raises:
            ValueError: If configuration is invalid
        """
        total = self.get_total_leds()
        if total <= 0:
            raise ValueError("Calibration must have at least one LED")

        for edge, count in [("top", self.leds_top), ("right", self.leds_right),
                            ("bottom", self.leds_bottom), ("left", self.leds_left)]:
            if count < 0:
                raise ValueError(f"LED count for {edge} must be non-negative, got {count}")

        for edge in ["top", "right", "bottom", "left"]:
            start, end = self.get_edge_span(edge)
            if not (0.0 <= start <= 1.0) or not (0.0 <= end <= 1.0):
                raise ValueError(f"Span for {edge} must be in [0.0, 1.0], got ({start}, {end})")
            if end <= start:
                raise ValueError(f"Span end must be greater than start for {edge}, got ({start}, {end})")

        return True

    def get_total_leds(self) -> int:
        """Get total number of LEDs across all edges."""
        return self.leds_top + self.leds_right + self.leds_bottom + self.leds_left

    def get_segment_for_edge(self, edge: str) -> CalibrationSegment | None:
        """Get segment configuration for a specific edge."""
        for seg in self.segments:
            if seg.edge == edge:
                return seg
        return None


class PixelMapper:
    """Maps screen border pixels to LED colors based on calibration."""

    def __init__(
        self,
        calibration: CalibrationConfig,
        interpolation_mode: Literal["average", "median", "dominant"] = "average",
    ):
        """Initialize pixel mapper.

        Args:
            calibration: Calibration configuration
            interpolation_mode: Color calculation mode
        """
        self.calibration = calibration
        self.interpolation_mode = interpolation_mode

        # Validate calibration
        self.calibration.validate()

        # Select color calculation function
        if interpolation_mode == "average":
            self._calc_color = calculate_average_color
        elif interpolation_mode == "median":
            self._calc_color = calculate_median_color
        elif interpolation_mode == "dominant":
            self._calc_color = calculate_dominant_color
        else:
            raise ValueError(f"Invalid interpolation mode: {interpolation_mode}")

        logger.info(
            f"Initialized pixel mapper with {self.calibration.get_total_leds()} LEDs "
            f"using {interpolation_mode} interpolation"
        )

    def map_border_to_leds(
        self,
        border_pixels: BorderPixels
    ) -> List[Tuple[int, int, int]]:
        """Map screen border pixels to LED colors.

        Args:
            border_pixels: Extracted border pixels from screen

        Returns:
            List of (R, G, B) tuples for each LED

        Raises:
            ValueError: If border pixels don't match calibration
        """
        total_leds = self.calibration.get_total_leds()
        led_colors = [(0, 0, 0)] * total_leds

        # Process each edge
        for edge_name in ["top", "right", "bottom", "left"]:
            segment = self.calibration.get_segment_for_edge(edge_name)

            if not segment:
                # This edge is not configured
                continue

            # Get pixels for this edge
            if edge_name == "top":
                edge_pixels = border_pixels.top
            elif edge_name == "right":
                edge_pixels = border_pixels.right
            elif edge_name == "bottom":
                edge_pixels = border_pixels.bottom
            else:  # left
                edge_pixels = border_pixels.left

            # Slice to span region if not full coverage
            span_start, span_end = self.calibration.get_edge_span(edge_name)
            if span_start > 0.0 or span_end < 1.0:
                if edge_name in ("top", "bottom"):
                    total_w = edge_pixels.shape[1]
                    s = int(span_start * total_w)
                    e = int(span_end * total_w)
                    edge_pixels = edge_pixels[:, s:e, :]
                else:
                    total_h = edge_pixels.shape[0]
                    s = int(span_start * total_h)
                    e = int(span_end * total_h)
                    edge_pixels = edge_pixels[s:e, :, :]

            # Divide edge into segments matching LED count
            try:
                pixel_segments = get_edge_segments(
                    edge_pixels,
                    segment.led_count,
                    edge_name
                )
            except ValueError as e:
                logger.error(f"Failed to segment {edge_name} edge: {e}")
                raise

            # Calculate LED indices for this segment
            led_indices = list(range(segment.led_start, segment.led_start + segment.led_count))

            # Reverse if needed
            if segment.reverse:
                led_indices = list(reversed(led_indices))

            # Map pixel segments to LEDs
            for led_idx, pixel_segment in zip(led_indices, pixel_segments):
                color = self._calc_color(pixel_segment)
                led_colors[led_idx] = color

        # Apply physical LED offset by rotating the array
        # Offset = number of LEDs from LED 0 to the start corner
        # Physical LED[i] should get calibration color[(i - offset) % total]
        offset = self.calibration.offset % total_leds if total_leds > 0 else 0
        if offset > 0:
            led_colors = led_colors[total_leds - offset:] + led_colors[:total_leds - offset]

        logger.debug(f"Mapped border pixels to {total_leds} LED colors (offset={offset})")
        return led_colors

    def test_calibration(self, edge: str, color: Tuple[int, int, int]) -> List[Tuple[int, int, int]]:
        """Generate test pattern to light up specific edge.

        Useful for verifying calibration configuration.

        Args:
            edge: Edge to light up (top, right, bottom, left)
            color: RGB color to use

        Returns:
            List of LED colors with only the specified edge lit

        Raises:
            ValueError: If edge is not in calibration
        """
        segment = self.calibration.get_segment_for_edge(edge)
        if not segment:
            raise ValueError(f"Edge '{edge}' not found in calibration")

        total_leds = self.calibration.get_total_leds()
        led_colors = [(0, 0, 0)] * total_leds

        # Light up the specified edge
        led_indices = range(segment.led_start, segment.led_start + segment.led_count)
        for led_idx in led_indices:
            led_colors[led_idx] = color

        logger.info(f"Generated test pattern for {edge} edge with color {color}")
        return led_colors


def create_default_calibration(led_count: int) -> CalibrationConfig:
    """Create a default calibration for a rectangular screen.

    Assumes LEDs are evenly distributed around the screen edges in clockwise order
    starting from bottom-left.

    Args:
        led_count: Total number of LEDs

    Returns:
        Default calibration configuration
    """
    if led_count < 4:
        raise ValueError("Need at least 4 LEDs for default calibration")

    # Distribute LEDs evenly across 4 edges
    leds_per_edge = led_count // 4
    remainder = led_count % 4

    # Distribute remainder to longer edges (bottom and top)
    bottom_count = leds_per_edge + (1 if remainder > 0 else 0)
    right_count = leds_per_edge
    top_count = leds_per_edge + (1 if remainder > 1 else 0)
    left_count = leds_per_edge + (1 if remainder > 2 else 0)

    config = CalibrationConfig(
        layout="clockwise",
        start_position="bottom_left",
        leds_bottom=bottom_count,
        leds_right=right_count,
        leds_top=top_count,
        leds_left=left_count,
    )

    logger.info(
        f"Created default calibration for {led_count} LEDs: "
        f"bottom={bottom_count}, right={right_count}, "
        f"top={top_count}, left={left_count}"
    )

    return config


def calibration_from_dict(data: dict) -> CalibrationConfig:
    """Create calibration configuration from dictionary.

    Supports both new format (leds_top/right/bottom/left) and legacy format
    (segments list) for backward compatibility.

    Args:
        data: Dictionary with calibration data

    Returns:
        CalibrationConfig instance

    Raises:
        ValueError: If data is invalid
    """
    try:
        config = CalibrationConfig(
            layout=data["layout"],
            start_position=data["start_position"],
            offset=data.get("offset", 0),
            leds_top=data.get("leds_top", 0),
            leds_right=data.get("leds_right", 0),
            leds_bottom=data.get("leds_bottom", 0),
            leds_left=data.get("leds_left", 0),
            span_top_start=data.get("span_top_start", 0.0),
            span_top_end=data.get("span_top_end", 1.0),
            span_right_start=data.get("span_right_start", 0.0),
            span_right_end=data.get("span_right_end", 1.0),
            span_bottom_start=data.get("span_bottom_start", 0.0),
            span_bottom_end=data.get("span_bottom_end", 1.0),
            span_left_start=data.get("span_left_start", 0.0),
            span_left_end=data.get("span_left_end", 1.0),
        )

        config.validate()
        return config

    except KeyError as e:
        raise ValueError(f"Missing required calibration field: {e}")
    except Exception as e:
        if isinstance(e, ValueError):
            raise
        raise ValueError(f"Invalid calibration data: {e}")


def calibration_to_dict(config: CalibrationConfig) -> dict:
    """Convert calibration configuration to dictionary.

    Args:
        config: Calibration configuration

    Returns:
        Dictionary representation
    """
    result = {
        "layout": config.layout,
        "start_position": config.start_position,
        "offset": config.offset,
        "leds_top": config.leds_top,
        "leds_right": config.leds_right,
        "leds_bottom": config.leds_bottom,
        "leds_left": config.leds_left,
    }
    # Include span fields only when not default (full coverage)
    for edge in ["top", "right", "bottom", "left"]:
        start = getattr(config, f"span_{edge}_start")
        end = getattr(config, f"span_{edge}_end")
        if start != 0.0 or end != 1.0:
            result[f"span_{edge}_start"] = start
            result[f"span_{edge}_end"] = end
    return result