"""Color correction postprocessing filter."""

import math
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


def _kelvin_to_rgb(kelvin: int) -> tuple:
    """Convert color temperature in Kelvin to normalized RGB multipliers.

    Uses Tanner Helland's approximation, normalized so 6500K = (1, 1, 1).
    """
    t = kelvin / 100.0

    # Red
    if t <= 66:
        r = 255.0
    else:
        r = 329.698727446 * ((t - 60) ** -0.1332047592)

    # Green
    if t <= 66:
        g = 99.4708025861 * math.log(t) - 161.1195681661
    else:
        g = 288.1221695283 * ((t - 60) ** -0.0755148492)

    # Blue
    if t >= 66:
        b = 255.0
    elif t <= 19:
        b = 0.0
    else:
        b = 138.5177312231 * math.log(t - 10) - 305.0447927307

    r = max(0.0, min(255.0, r))
    g = max(0.0, min(255.0, g))
    b = max(0.0, min(255.0, b))

    return r / 255.0, g / 255.0, b / 255.0


# Pre-compute 6500K reference for normalization
_REF_R, _REF_G, _REF_B = _kelvin_to_rgb(6500)


@FilterRegistry.register
class ColorCorrectionFilter(PostprocessingFilter):
    """Adjusts color temperature and per-channel RGB gains using LUTs."""

    filter_id = "color_correction"
    filter_name = "Color Correction"

    def __init__(self, options: Dict[str, Any]):
        super().__init__(options)
        temp = self.options["temperature"]
        rg = self.options["red_gain"]
        gg = self.options["green_gain"]
        bg = self.options["blue_gain"]

        # Color temperature -> RGB multipliers, normalized to 6500K = (1,1,1)
        tr, tg, tb = _kelvin_to_rgb(temp)
        r_mult = (tr / _REF_R) * rg
        g_mult = (tg / _REF_G) * gg
        b_mult = (tb / _REF_B) * bg

        # Build merged (256, 3) LUT for single-pass numpy fancy-index lookup
        src = np.arange(256, dtype=np.float32)
        lut_r = np.clip(src * r_mult, 0, 255).astype(np.uint8)
        lut_g = np.clip(src * g_mult, 0, 255).astype(np.uint8)
        lut_b = np.clip(src * b_mult, 0, 255).astype(np.uint8)
        self._lut = np.stack([lut_r, lut_g, lut_b], axis=-1)  # (256, 3)

        self._is_neutral = (temp == 6500 and rg == 1.0 and gg == 1.0 and bg == 1.0)

    @classmethod
    def get_options_schema(cls) -> List[FilterOptionDef]:
        return [
            FilterOptionDef(
                key="temperature",
                label="Color Temperature (K)",
                option_type="int",
                default=6500,
                min_value=2000,
                max_value=10000,
                step=100,
            ),
            FilterOptionDef(
                key="red_gain",
                label="Red Gain",
                option_type="float",
                default=1.0,
                min_value=0.0,
                max_value=2.0,
                step=0.05,
            ),
            FilterOptionDef(
                key="green_gain",
                label="Green Gain",
                option_type="float",
                default=1.0,
                min_value=0.0,
                max_value=2.0,
                step=0.05,
            ),
            FilterOptionDef(
                key="blue_gain",
                label="Blue Gain",
                option_type="float",
                default=1.0,
                min_value=0.0,
                max_value=2.0,
                step=0.05,
            ),
        ]

    def process_image(self, image: np.ndarray, image_pool: ImagePool) -> Optional[np.ndarray]:
        if self._is_neutral:
            return None
        image[:] = self._lut[image]
        return None
