"""
Unit tests for SunburstChartStrategy._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.sunburst import SunburstChartStrategy
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_sunburst(**kwargs: Any) -> SunburstChartStrategy:
    return SunburstChartStrategy(
        fields_metadata={"path": _cat_meta()},
        path_variables=["path"],
        aggregation=Aggregation(name="count", function="count"),
        **kwargs,
    )


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


def test_prepare_plot_data_returns_categories_values_colors(
    sunburst_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_sunburst()
    result = instance._prepare_plot_data({"main": sunburst_df}, tr, theme)
    assert len(result["categories"]) == 3
    assert len(result["values"]) == 3
    assert len(result["colors"]) == 3
    assert all(c.startswith("#") for c in result["colors"])


def test_prepare_plot_data_sort_by_value_descending(
    sunburst_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_sunburst(sort_by_value=SortOrder.DESCENDING)
    result = instance._prepare_plot_data({"main": sunburst_df}, tr, theme)
    values = result["values"]
    assert values == sorted(values, reverse=True)


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