"""
Unit tests for ChoroplethMapChartStrategy._prepare_plot_data() and _render_from_payload().
"""

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.choropleth_map import ChoroplethMapChartStrategy
from orama.variables import Aggregation


def _country_meta() -> MagicMock:
    m = MagicMock(spec=FieldMetadata)
    m.categorical = False
    m.numeric = False
    m.derived = False
    m.effective_type = str
    m.extra = {"suggested_validation": "iso_a2"}
    return m


def _make_choropleth(**kwargs: Any) -> ChoroplethMapChartStrategy:
    return ChoroplethMapChartStrategy(
        fields_metadata={"country": _country_meta()},
        country_variable="country",
        aggregation=Aggregation(name="count", function="count"),
        **kwargs,
    )


@pytest.fixture
def choropleth_df() -> pl.DataFrame:
    return pl.DataFrame({"country": ["US", "DE"], "count": [100, 200]})


def test_prepare_plot_data_returns_country_codes_names_values(
    choropleth_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_choropleth()
    result = instance._prepare_plot_data({"main": choropleth_df}, tr, theme)
    # sorted alphabetically by country code
    assert result["country_codes"] == ["DE", "US"]
    assert len(result["country_names"]) == 2
    assert "Germany" in result["country_names"]
    assert "United States" in result["country_names"]
    assert len(result["values"]) == 2


def test_prepare_plot_data_continuous_mode_sets_color_scale(
    choropleth_df: pl.DataFrame,
    tr: Any,
    theme: Any,
) -> None:
    instance = _make_choropleth()
    result = instance._prepare_plot_data({"main": choropleth_df}, tr, theme)
    assert result["_use_discrete"] is False
    assert result["_color_scale"] is not None
    assert result["_discrete_colors"] is None


def test_prepare_plot_data_discrete_mode_with_categorical_rule(
    choropleth_df: pl.DataFrame,
    anchor_colors: list[HSLColor],
    tr: Any,
    theme: Any,
) -> None:
    rule = ColorRule(
        assigner=ColorAssignerByCategory(anchor_colors=anchor_colors),
        selector=ColorSelectorStrategy(),
    )
    instance = _make_choropleth(color_rules=[rule])
    result = instance._prepare_plot_data({"main": choropleth_df}, tr, theme)
    assert result["_use_discrete"] is True
    assert isinstance(result["_discrete_colors"], list)
    assert len(result["_discrete_colors"]) == 2
    assert all(c.startswith("#") for c in result["_discrete_colors"])


def _minimal_payload() -> dict[str, Any]:
    return {
        "country_names": ["Germany"],
        "values": [100],
        "_alpha3_codes": ["DEU"],
        "_use_discrete": True,
        "_discrete_colors": ["#ff0000"],
        "_color_scale": None,
        "title": "Test",
    }


@pytest.mark.parametrize("resolution_high,expected", [(True, 50), (False, 110)])
def test_render_resolution(resolution_high: bool, expected: int) -> None:
    instance = _make_choropleth(resolution_high=resolution_high)
    fig = instance._render_from_payload(_minimal_payload())
    assert fig.layout.geo.resolution == expected


@pytest.mark.parametrize("show_countries", [True, False])
def test_render_show_countries(show_countries: bool) -> None:
    instance = _make_choropleth(show_countries=show_countries)
    fig = instance._render_from_payload(_minimal_payload())
    assert fig.layout.geo.showcountries is show_countries
    assert fig.layout.geo.showlakes is show_countries


@pytest.mark.parametrize("show_ocean", [True, False])
def test_render_show_ocean(show_ocean: bool) -> None:
    instance = _make_choropleth(show_ocean=show_ocean)
    fig = instance._render_from_payload(_minimal_payload())
    assert fig.layout.geo.showocean is show_ocean
