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

from typing import Any
from unittest.mock import MagicMock

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

from orama.color import (
    ColorAssignerByCategory,
    ColorRule,
    ColorSelectorStrategy,
)
from orama.enums import SortOrder
from orama.strategies.pie import PieChartStrategy
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_pie(**kwargs: Any) -> PieChartStrategy:
    return PieChartStrategy(
        fields_metadata={"cat": _cat_meta()},
        category_variables=["cat"],
        aggregation=Aggregation(name="count", function="count"),
        **kwargs,
    )


@pytest.fixture
def pie_df() -> pl.DataFrame:
    return pl.DataFrame({"cat": ["A", "B", "C"], "count": [10, 20, 30]})


def test_prepare_plot_data_returns_categories_values_colors(
    pie_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    pie = _make_pie()
    result = pie._prepare_plot_data({"main": pie_df}, tr, theme)
    assert result["categories"] == ["A", "B", "C"]
    assert result["values"] == [10, 20, 30]
    assert len(result["colors"]) == 3
    assert all(isinstance(c, str) and c.startswith("#") for c in result["colors"])


def test_prepare_plot_data_preserves_category_order(
    tr: Any,
    theme: Any,
) -> None:
    pie = _make_pie()
    df = pl.DataFrame({"cat": ["Z", "A", "M"], "count": [1, 2, 3]})
    result = pie._prepare_plot_data({"main": df}, tr, theme)
    assert result["categories"] == ["A", "M", "Z"]


def test_prepare_plot_data_sort_by_value_descending(
    tr: Any,
    theme: Any,
) -> None:
    pie = _make_pie(sort_by_value=SortOrder.DESCENDING)
    df = pl.DataFrame({"cat": ["A", "B", "C"], "count": [10, 30, 20]})
    result = pie._prepare_plot_data({"main": df}, tr, theme)
    values = result["values"]
    assert values == sorted(values, reverse=True)


def test_prepare_plot_data_color_rules_applied(
    pie_df: pl.DataFrame,
    anchor_colors: list[HSLColor],
    tr: Any,
    theme: Any,
) -> None:
    rule = ColorRule(
        assigner=ColorAssignerByCategory(anchor_colors=anchor_colors),
        selector=ColorSelectorStrategy(),
    )
    pie = _make_pie(color_rules=[rule])
    result = pie._prepare_plot_data({"main": pie_df}, tr, theme)
    assert len(result["colors"]) == 3
    assert all(c.startswith("#") for c in result["colors"])
