"""Palette quantization postprocessing filter."""

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

import numpy as np

from wled_controller.core.filters.base import FilterOptionDef, PostprocessingFilter
from wled_controller.core.filters.image_pool import ImagePool
from wled_controller.core.filters.registry import FilterRegistry

_PRESETS = {
    "warm": "#FF4500,#FF6347,#FF8C00,#FFD700,#FFA500,#DC143C",
    "cool": "#0000FF,#1E90FF,#00BFFF,#00CED1,#4682B4,#6A5ACD",
    "neon": "#FF00FF,#00FFFF,#00FF00,#FFFF00,#FF0000,#0000FF",
    "pastel": "#FFB3BA,#FFDFBA,#FFFFBA,#BAFFC9,#BAE1FF,#E8BAFF",
    "monochrome": "#000000,#404040,#808080,#C0C0C0,#FFFFFF",
    "rgb": "#FF0000,#00FF00,#0000FF",
    "rainbow": "#FF0000,#FF7F00,#FFFF00,#00FF00,#0000FF,#4B0082,#8F00FF",
    "earth": "#8B4513,#A0522D,#D2691E,#228B22,#006400,#2F4F4F",
}

_HEX_RE = re.compile(r"^#?([0-9a-fA-F]{6})$")


def _parse_hex_colors(color_str: str) -> Optional[np.ndarray]:
    """Parse comma-separated hex colors into (N, 3) uint8 array."""
    colors = []
    for part in color_str.split(","):
        part = part.strip()
        m = _HEX_RE.match(part)
        if m:
            h = m.group(1)
            colors.append((int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)))
    if not colors:
        return None
    return np.array(colors, dtype=np.uint8)


@FilterRegistry.register
class PaletteQuantizationFilter(PostprocessingFilter):
    """Maps every pixel to the nearest color in a palette."""

    filter_id = "palette_quantization"
    filter_name = "Palette Quantization"

    def __init__(self, options: Dict[str, Any]):
        super().__init__(options)
        self._palette: Optional[np.ndarray] = None
        self._rebuild_palette()

    def _rebuild_palette(self):
        mode = self.options.get("mode", "preset")
        if mode == "custom":
            color_str = self.options.get("colors", "")
            self._palette = _parse_hex_colors(color_str)
        else:
            preset_name = self.options.get("preset", "rainbow")
            preset_str = _PRESETS.get(preset_name, _PRESETS["rainbow"])
            self._palette = _parse_hex_colors(preset_str)

    @classmethod
    def get_options_schema(cls) -> List[FilterOptionDef]:
        return [
            FilterOptionDef(
                key="mode",
                label="Mode",
                option_type="select",
                default="preset",
                min_value=None,
                max_value=None,
                step=None,
                choices=[
                    {"value": "preset", "label": "Preset"},
                    {"value": "custom", "label": "Custom"},
                ],
            ),
            FilterOptionDef(
                key="preset",
                label="Preset Palette",
                option_type="select",
                default="rainbow",
                min_value=None,
                max_value=None,
                step=None,
                choices=[{"value": k, "label": k.capitalize(), "colors": v} for k, v in _PRESETS.items()],
            ),
            FilterOptionDef(
                key="colors",
                label="Custom Colors (hex, comma-separated)",
                option_type="string",
                default="#FF0000,#00FF00,#0000FF",
                min_value=None,
                max_value=None,
                step=None,
                max_length=500,
            ),
        ]

    def process_image(self, image: np.ndarray, image_pool: ImagePool) -> Optional[np.ndarray]:
        if self._palette is None or len(self._palette) == 0:
            return None

        palette = self._palette.astype(np.int16)
        flat = image.reshape(-1, 3)
        total = flat.shape[0]
        chunk_size = 32768

        for start in range(0, total, chunk_size):
            end = min(start + chunk_size, total)
            chunk = flat[start:end].astype(np.int16)
            # Squared Euclidean distance: (C, 1, 3) - (1, N, 3) -> (C, N, 3) -> sum -> (C, N)
            diff = chunk[:, np.newaxis, :] - palette[np.newaxis, :, :]
            dist_sq = (diff * diff).sum(axis=2)
            nearest = dist_sq.argmin(axis=1)
            flat[start:end] = self._palette[nearest]

        return None
