"""
Unit tests for GroupedBarChartStrategy helper methods.
"""

from typing import Any, cast
from unittest.mock import MagicMock

import polars as pl
import pytest
from fields_metadata import FieldMetadata
from polychromos import HSLColor

from orama.color import (
    ColorAssignerByGroup,
    ColorRule,
    ColorSelectorStrategy,
)
from orama.strategies.grouped_bar import GroupedBarChartStrategy
from orama.variables import Aggregation


def _cat_meta() -> MagicMock:
    m = MagicMock(spec=FieldMetadata)
    m.categorical = True
    m.numeric = False
    m.derived = False
    m.effective_type = str
    m.extra = {}
    return m


def _make_grouped_bar(**kwargs: Any) -> GroupedBarChartStrategy:
    fields = cast(dict[str, FieldMetadata], {"cat": _cat_meta(), "grp": _cat_meta()})
    return GroupedBarChartStrategy(
        fields_metadata=fields,
        category_variables=["cat"],
        group_variable="grp",
        aggregation=Aggregation(name="count", function="count"),
        **kwargs,
    )


@pytest.fixture
def base_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "cat": ["A", "A", "B", "B"],
            "grp": ["X", "Y", "X", "Y"],
            "count": [10, 20, 30, 40],
        }
    )


def test_prepare_grouped_df_returns_correct_columns(
    base_df: pl.DataFrame,
    tr: Any,
) -> None:
    instance = _make_grouped_bar()
    df, alias, value_var, sorted_groups = instance._prepare_grouped_df(base_df, tr)
    assert alias in df.columns
    assert "grp" in df.columns
    assert "count" in df.columns
    assert "category_total" in df.columns


def test_prepare_grouped_df_relative_adds_new_value_col(
    base_df: pl.DataFrame,
    tr: Any,
) -> None:
    instance = _make_grouped_bar(stacked=True, relative=True)
    df, alias, value_var, sorted_groups = instance._prepare_grouped_df(base_df, tr)
    assert value_var == "Relative count"
    assert "Relative count" in df.columns


def test_compute_group_color_map_no_rules_returns_none(
    theme: Any,
) -> None:
    instance = _make_grouped_bar()
    result = instance._compute_group_color_map(["X", "Y"], theme)
    assert result is None


def test_compute_group_color_map_with_rules_returns_dict(
    anchor_colors: list[HSLColor],
    theme: Any,
) -> None:
    assigner = ColorAssignerByGroup(anchor_colors=anchor_colors)
    rule = ColorRule(assigner=assigner, selector=ColorSelectorStrategy())
    instance = _make_grouped_bar(color_rules=[rule])
    result = instance._compute_group_color_map(["X", "Y"], theme)
    assert isinstance(result, dict)
    assert set(result.keys()) == {"X", "Y"}
    for v in result.values():
        assert v.startswith("#")


def test_prepare_plot_data_returns_categories_groups_color_map(
    base_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_grouped_bar()
    result = instance._prepare_plot_data({"main": base_df}, tr, theme)
    assert isinstance(result["categories"], list)
    assert isinstance(result["groups"], list)
    assert result["color_map"] is None


def test_prepare_plot_data_color_map_with_rules(
    base_df: pl.DataFrame,
    anchor_colors: list[HSLColor],
    tr: Any,
    theme: Any,
) -> None:
    assigner = ColorAssignerByGroup(anchor_colors=anchor_colors)
    rule = ColorRule(assigner=assigner, selector=ColorSelectorStrategy())
    instance = _make_grouped_bar(color_rules=[rule])
    result = instance._prepare_plot_data({"main": base_df}, tr, theme)
    assert isinstance(result["color_map"], dict)
    assert set(result["color_map"].keys()) == {"X", "Y"}
    for v in result["color_map"].values():
        assert v.startswith("#")
