import polars as pl
import pytest
from polychromos import HSLColor

from orama.color import (
    ColorAssignerByCategory,
    ColorAssignerByGroup,
    ColorAssignerByValue,
    ColorRule,
    ColorSelectByCategoryName,
    ColorSelectMaxValue,
    ColorSelectMinValue,
    ColorSelectorStrategy,
    ColorSelectValuesAbove,
    ColorSelectValuesAboveMean,
    ColorSelectValuesBelow,
    ColorSelectValuesBelowMean,
)

_DF = pl.DataFrame({"cat": ["A", "B", "C"], "val": [10.0, 50.0, 30.0]})
_DF_MEAN = pl.DataFrame(
    {
        "cat": ["A", "B", "C"],
        "val": [10.0, 50.0, 30.0],
        "__total_mean__": [30.0, 30.0, 30.0],
    }
)
_EMPTY_DF = pl.DataFrame({"cat": [], "val": []})


@pytest.mark.parametrize(
    "selector,expected",
    [
        (ColorSelectorStrategy(), [True, True, True]),
        (ColorSelectByCategoryName(["A", "C"]), [True, False, True]),
        (ColorSelectValuesBelow(35.0), [True, False, True]),
        (ColorSelectValuesAbove(25.0), [False, True, True]),
        (ColorSelectMaxValue(), [False, True, False]),
        (ColorSelectMinValue(), [True, False, False]),
    ],
)
def test_color_selector_returns_expected(selector: ColorSelectorStrategy, expected: list[bool]) -> None:
    result = selector.selector(_DF, "cat", "val")
    assert result == expected


def test_color_select_max_value_empty_df() -> None:
    result = ColorSelectMaxValue().selector(_EMPTY_DF, "cat", "val")
    assert result == []


def test_color_select_min_value_empty_df() -> None:
    result = ColorSelectMinValue().selector(_EMPTY_DF, "cat", "val")
    assert result == []


def test_color_select_values_below_mean() -> None:
    result = ColorSelectValuesBelowMean().selector(_DF_MEAN, "cat", "val")
    assert result == [True, False, True]


def test_color_select_values_above_mean() -> None:
    result = ColorSelectValuesAboveMean().selector(_DF_MEAN, "cat", "val")
    assert result == [False, True, True]


def test_color_assigner_strategy_empty_anchor_colors_raises() -> None:
    with pytest.raises(ValueError):
        ColorAssignerByCategory(anchor_colors=[])


def test_color_assigner_strategy_anchor_colors_property(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    assert assigner.anchor_colors == anchor_colors


def test_color_assigner_strategy_continuous_is_false(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    assert assigner.continuous is False


def test_color_assigner_by_category_assign_colors(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    colors = assigner.assign_colors(_DF, "cat", "val")
    assert len(colors) == len(_DF["cat"])


def test_color_assigner_by_category_continuous(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    assert assigner.continuous is False


def test_color_assigner_by_group_no_group_var_raises(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByGroup(anchor_colors=anchor_colors)
    with pytest.raises(ValueError):
        assigner.assign_colors(_DF, "cat", "val", group_var=None)


def test_color_assigner_by_group_assign_colors(anchor_colors: list[HSLColor]) -> None:
    df = pl.DataFrame({"cat": ["A", "A", "B"], "val": [1.0, 2.0, 3.0], "group": ["G1", "G1", "G2"]})
    assigner = ColorAssignerByGroup(anchor_colors=anchor_colors)
    colors = assigner.assign_colors(df, "cat", "val", group_var="group")
    assert len(colors) == len(df["group"].unique())


def test_color_assigner_by_value_continuous(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByValue(anchor_colors=anchor_colors)
    assert assigner.continuous is True


def test_color_assigner_by_value_properties(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByValue(
        anchor_colors=anchor_colors,
        range_min=0.0,
        range_max=100.0,
        range_center_at=50.0,
    )
    assert assigner.range_min == 0.0
    assert assigner.range_max == 100.0
    assert assigner.range_center_at == 50.0


def test_color_assigner_by_value_properties_none_defaults(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByValue(anchor_colors=anchor_colors)
    assert assigner.range_min is None
    assert assigner.range_max is None
    assert assigner.range_center_at is None


def test_color_assigner_by_value_assign_colors_length(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByValue(anchor_colors=anchor_colors)
    colors = assigner.assign_colors(_DF, "cat", "val")
    assert len(colors) == _DF.height


def test_color_rule_default_selector(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    rule = ColorRule(assigner=assigner)
    assert isinstance(rule.selector, ColorSelectorStrategy)


def test_color_rule_properties(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    selector = ColorSelectByCategoryName(["A"])
    rule = ColorRule(assigner=assigner, selector=selector)
    assert rule.selector is selector
    assert rule.assigner is assigner


@pytest.mark.parametrize(
    "selector,expected",
    [
        (ColorSelectorStrategy(), "selecting all entries"),
        (ColorSelectByCategoryName(["A", "B"]), "selecting entries in categories ['A', 'B']"),
        (ColorSelectValuesBelow(10.0), "selecting entries with a value below 10.0"),
        (ColorSelectValuesAbove(5.0), "selecting entries with a value above 5.0"),
        (ColorSelectMaxValue(), "selecting the entry with the highest value"),
        (ColorSelectMinValue(), "selecting the entry with the lowest value"),
        (ColorSelectValuesBelowMean(), "selecting entries with a value below the mean"),
        (ColorSelectValuesAboveMean(), "selecting entries with a value above the mean"),
    ],
)
def test_color_selector_explain(selector: ColorSelectorStrategy, expected: str) -> None:
    assert selector.explain() == expected


def test_color_assigner_by_category_explain(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    result = assigner.explain()
    assert result.startswith("colorize them each in a distinct color from")
    assert anchor_colors[0].to_css_hex() in result


def test_color_assigner_by_group_explain(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByGroup(anchor_colors=anchor_colors)
    result = assigner.explain()
    assert result.startswith("colorize them each in a distinct color by group from")
    assert anchor_colors[0].to_css_hex() in result


def test_color_assigner_by_value_explain_no_range(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByValue(anchor_colors=anchor_colors)
    result = assigner.explain()
    assert result.startswith("colorize them by their value in colors from")
    assert anchor_colors[0].to_css_hex() in result
    assert "restricted" not in result
    assert "centered" not in result


def test_color_assigner_by_value_explain_with_range(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByValue(anchor_colors=anchor_colors, range_min=0.0, range_max=100.0)
    result = assigner.explain()
    assert "restricted to values 0.0\u2013100.0" in result


def test_color_assigner_by_value_explain_with_center(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByValue(anchor_colors=anchor_colors, range_center_at=50.0)
    result = assigner.explain()
    assert "centered at 50.0" in result


def test_color_rule_explain(anchor_colors: list[HSLColor]) -> None:
    selector = ColorSelectValuesAbove(12.0)
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    rule = ColorRule(assigner=assigner, selector=selector)
    result = rule.explain()
    assert result.startswith("selecting entries with a value above 12.0")
    assert "colorize them" in result


def test_color_rule_explain_default_selector(anchor_colors: list[HSLColor]) -> None:
    assigner = ColorAssignerByCategory(anchor_colors=anchor_colors)
    rule = ColorRule(assigner=assigner)
    result = rule.explain()
    assert result.startswith("selecting all entries")
    assert "colorize them" in result
