"""
Unit tests for GroupedBoxplotChartStrategy._prepare_plot_data().
"""

from typing import Any
from unittest.mock import MagicMock

import polars as pl
import pytest
from fields_metadata import FieldMetadata

from orama.strategies.grouped_boxplot import GroupedBoxplotChartStrategy


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 _num_meta() -> MagicMock:
    m = MagicMock(spec=FieldMetadata)
    m.categorical = False
    m.numeric = True
    m.derived = False
    m.effective_type = float
    m.extra = {}
    return m


def _make_grouped_boxplot(**kwargs: Any) -> GroupedBoxplotChartStrategy:
    return GroupedBoxplotChartStrategy(
        fields_metadata={"cat": _cat_meta(), "val": _num_meta(), "grp": _cat_meta()},
        category_variables=["cat"],
        value_variable="val",
        group_variable="grp",
        **kwargs,
    )


@pytest.fixture
def grouped_stats_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "cat": ["A", "A", "B", "B"],
            "grp": ["X", "Y", "X", "Y"],
            "q1": [1.0, 1.5, 2.0, 2.5],
            "median": [2.0, 2.5, 3.0, 3.5],
            "q3": [3.0, 3.5, 4.0, 4.5],
            "mean": [2.0, 2.5, 3.0, 3.5],
            "sd": [0.5, 0.5, 0.5, 0.5],
            "min": [0.0, 0.5, 1.0, 1.5],
            "max": [4.0, 4.5, 5.0, 5.5],
        }
    )


def test_prepare_plot_data_returns_categories_groups_color_map(
    grouped_stats_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_grouped_boxplot()
    result = instance._prepare_plot_data({"main": grouped_stats_df}, tr, theme)
    assert isinstance(result["categories"], list)
    assert isinstance(result["groups"], list)
    assert isinstance(result["color_map"], dict)
    assert "_group_colors" in result


def test_prepare_plot_data_color_map_keyed_by_group_values(
    grouped_stats_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_grouped_boxplot()
    result = instance._prepare_plot_data({"main": grouped_stats_df}, tr, theme)
    assert set(result["color_map"].keys()) == {"X", "Y"}
    for v in result["color_map"].values():
        assert v.startswith("#")
