"""
Unit tests for CategoricalLineChartStrategy._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.strategies.categorical_line import CategoricalLineChartStrategy
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_cat_line(extra_fields: dict | None = None, **kwargs: Any) -> CategoricalLineChartStrategy:
    fields = {"x": _cat_meta()}
    if extra_fields:
        fields.update(extra_fields)
    return CategoricalLineChartStrategy(
        fields_metadata=fields,
        x_variables=["x"],
        aggregation=Aggregation(name="count", function="count"),
        **kwargs,
    )


@pytest.fixture
def x_df() -> pl.DataFrame:
    return pl.DataFrame({"x": ["Q1", "Q2", "Q3"], "count": [10, 20, 30]})


def test_prepare_plot_data_no_series_returns_x_labels_and_one_color(
    x_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_cat_line()
    result = instance._prepare_plot_data({"main": x_df}, tr, theme)
    assert isinstance(result["x_labels"], list)
    assert len(result["x_labels"]) == 3
    assert result["series"] == [None]
    assert len(result["colors"]) == 1
    assert result["colors"][0].startswith("#")


def test_prepare_plot_data_with_series_variable(
    tr: Any,
    theme: Any,
) -> None:
    df = pl.DataFrame(
        {
            "x": ["Q1", "Q1", "Q2", "Q2"],
            "ser": ["A", "B", "A", "B"],
            "count": [10, 20, 30, 40],
        }
    )
    instance = _make_cat_line(extra_fields={"ser": _cat_meta()}, series_variable="ser")
    result = instance._prepare_plot_data({"main": df}, tr, theme)
    assert set(result["series"]) == {"A", "B"}
    assert len(result["colors"]) == 2
    assert all(c.startswith("#") for c in result["colors"])


def test_prepare_plot_data_color_rules_applied(
    x_df: pl.DataFrame,
    anchor_colors: list[HSLColor],
    tr: Any,
    theme: Any,
) -> None:
    extra_fields = {"ser": _cat_meta()}
    df = pl.DataFrame(
        {
            "x": ["Q1", "Q1"],
            "ser": ["A", "B"],
            "count": [5, 10],
        }
    )
    rule = ColorRule(
        assigner=ColorAssignerByCategory(anchor_colors=anchor_colors),
        selector=ColorSelectorStrategy(),
    )
    instance = _make_cat_line(extra_fields=extra_fields, series_variable="ser", color_rules=[rule])
    result = instance._prepare_plot_data({"main": df}, tr, theme)
    assert len(result["colors"]) == 2
    assert all(c.startswith("#") for c in result["colors"])
