"""
Unit tests for HeatMapChartStrategy._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.heatmap import HeatMapChartStrategy
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_heatmap(**kwargs: Any) -> HeatMapChartStrategy:
    return HeatMapChartStrategy(
        fields_metadata={"x": _cat_meta(), "y": _cat_meta()},
        x_category_variables=["x"],
        y_category_variables=["y"],
        aggregation=Aggregation(name="count", function="count"),
        **kwargs,
    )


@pytest.fixture
def heatmap_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "x": ["A", "A", "B", "B"],
            "y": ["P", "Q", "P", "Q"],
            "count": [1, 2, 3, 4],
        }
    )


@pytest.fixture
def two_color_theme(theme: Any, anchor_color: Any) -> Any:
    """Theme mock with two anchor colors so Palette.to_color_scale doesn't raise."""
    from polychromos import HSLColor

    second = HSLColor(hue=0.5, saturation=0.5, lightness=0.5)
    theme.get_default_color_sequence.return_value = [anchor_color, second]
    return theme


def test_prepare_plot_data_returns_x_y_categories_and_values(
    heatmap_df: pl.DataFrame,
    tr: Any,
    two_color_theme: Any,
) -> None:
    heatmap = _make_heatmap()
    result = heatmap._prepare_plot_data({"main": heatmap_df}, tr, two_color_theme)
    assert len(result["x_categories"]) == 4
    assert len(result["y_categories"]) == 4
    assert len(result["values"]) == 4
    assert set(result["values"]) == {1, 2, 3, 4}


def test_prepare_plot_data_color_scale_present(
    heatmap_df: pl.DataFrame,
    tr: Any,
    two_color_theme: Any,
) -> None:
    heatmap = _make_heatmap()
    result = heatmap._prepare_plot_data({"main": heatmap_df}, tr, two_color_theme)
    assert result["_color_scale"] is not None
