"""
Choropleth tile map figure strategy implementation.
"""

from typing import Any

import plotly.graph_objects as go
import polars as pl
import pycountry
from babel.support import Translations
from fields_metadata import FieldMetadata
from polychromos.color import HSLColor
from polychromos.palette import HSLColorScale, HSLColorSequence, Palette
from therismos.grouping import Aggregation as TherismosAggregation
from therismos.grouping import AggregationFunction

from orama.color import (
    ColorAssignerByCategory,
    ColorAssignerByValue,
    ColorAssignerStrategy,
    ColorRule,
    ColorSelectByCategoryName,
    ColorSelectMaxValue,
    ColorSelectMinValue,
    ColorSelectorStrategy,
    ColorSelectValuesAbove,
    ColorSelectValuesAboveMean,
    ColorSelectValuesBelow,
    ColorSelectValuesBelowMean,
)
from orama.enums import GeoProjection
from orama.options import (
    BooleanOption,
    ChartOptionsMap,
    GeoProjectionOption,
)
from orama.query import QueryParams
from orama.strategies.base import FigureStrategy
from orama.theme import Theme
from orama.variables import (
    Aggregation,
    ChartVariablesMap,
    CountryCodeVariable,
    NumericalAggregationVariable,
)


class ChoroplethMapChartStrategy(FigureStrategy):
    """
    Choropleth tile map figure strategy.

    Renders country-level data on an interactive world map using ISO 3166-1 alpha-2
    country codes. No Mapbox token is required.
    """

    _variables: ChartVariablesMap = {
        "country": CountryCodeVariable(
            name="Country",
            description=(
                "A categorical variable containing ISO 3166-1 alpha-2 country codes. "
                "Each row maps one country code to an aggregated value."
            ),
            map_to="country_variable",
        ),
        "value": NumericalAggregationVariable(
            name="Value",
            description=("An aggregation of a numerical variable to produce the value displayed for each country on the map."),
            map_to="aggregation",
        ),
    }
    _options: ChartOptionsMap = {
        "map_style": GeoProjectionOption(
            name="Map projection",
            description="The geo projection type to use for rendering the choropleth map.",
            default_value=GeoProjection.EQUIRECTANGULAR,
            map_to="map_style",
        ),
        "fit_bounds": BooleanOption(
            name="Auto-fit bounds",
            description="Automatically zoom the map to the extent of the countries with data.",
            default_value=False,
            map_to="fit_bounds",
        ),
        "resolution_high": BooleanOption(
            name="High resolution",
            description="Use high resolution (50) for map geo features; low resolution (110) otherwise.",
            default_value=True,
            map_to="resolution_high",
        ),
        "show_countries": BooleanOption(
            name="Show countries",
            description="Show country borders and lakes. When disabled, only coastlines and ocean are shown.",
            default_value=True,
            map_to="show_countries",
        ),
        "show_ocean": BooleanOption(
            name="Show ocean",
            description="Show ocean fill on the map.",
            default_value=True,
            map_to="show_ocean",
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        country_variable: str,
        aggregation: Aggregation,
        map_style: GeoProjection = GeoProjection.EQUIRECTANGULAR,
        fit_bounds: bool = False,
        resolution_high: bool = True,
        show_countries: bool = True,
        show_ocean: bool = True,
        color_rules: list[ColorRule] | None = None,
        title: str | None = None,
        subtitle: str | None = None,
        caption: str | None = None,
        description: str | None = None,
        storytelling_context: str | None = None,
        country_label: str | None = None,
        value_label: str | None = None,
    ) -> None:
        """
        Constructor method.

        :param fields_metadata: Mapping of field paths to their metadata.
        :type fields_metadata: dict[str, FieldMetadata]
        :param country_variable: Field containing ISO 3166-1 alpha-2 country codes.
        :type country_variable: str
        :param aggregation: Aggregation applied to the value field.
        :type aggregation: Aggregation
        :param map_style: Geo projection type. Defaults to ``GeoProjection.EQUIRECTANGULAR``.
        :type map_style: GeoProjection, optional
        :param fit_bounds: Automatically zoom the map to the extent of the countries with data.
            Defaults to ``False``.
        :type fit_bounds: bool, optional
        :param resolution_high: Use high resolution (50) for geo features; low resolution (110)
            otherwise. Defaults to ``True``.
        :type resolution_high: bool, optional
        :param show_countries: Show country borders and lakes. When ``False``, only coastlines
            and ocean are shown. Defaults to ``True``.
        :type show_countries: bool, optional
        :param show_ocean: Show ocean fill on the map. Defaults to ``True``.
        :type show_ocean: bool, optional
        :param color_rules: Optional list of color rules. Defaults to ``None``.
        :type color_rules: list[ColorRule] | None, optional
        :param title: Optional chart title. Defaults to ``None``.
        :type title: str | None, optional
        :param subtitle: Optional chart subtitle. Defaults to ``None``.
        :type subtitle: str | None, optional
        :param caption: Optional chart caption. Defaults to ``None``.
        :type caption: str | None, optional
        :param description: Optional free-form text explaining how to interpret the figure.
            Defaults to ``None``.
        :type description: str | None, optional
        :param country_label: Optional label for the country axis. Defaults to ``None``.
        :type country_label: str | None, optional
        :param value_label: Optional label for the value axis. Defaults to ``None``.
        :type value_label: str | None, optional
        """
        super().__init__(
            fields_metadata=fields_metadata,
            color_rules=color_rules,
            title=title,
            subtitle=subtitle,
            caption=caption,
            description=description,
            storytelling_context=storytelling_context,
        )

        self._variables["country"].validate_fields(fields_metadata, country_variable)
        if aggregation.field:
            self._variables["value"].validate_fields(fields_metadata, aggregation.field)

        self._options["map_style"].validate(map_style)
        self._options["fit_bounds"].validate(fit_bounds)
        self._options["resolution_high"].validate(resolution_high)
        self._options["show_countries"].validate(show_countries)
        self._options["show_ocean"].validate(show_ocean)

        self._country_variable = country_variable
        self._aggregation = aggregation
        self._map_style = map_style
        self._fit_bounds = fit_bounds
        self._resolution_high = resolution_high
        self._show_countries = show_countries
        self._show_ocean = show_ocean
        self._country_label = country_label
        self._value_label = value_label

    @classmethod
    def name(cls) -> str:
        return "Choropleth Map"

    @classmethod
    def description(cls) -> str:
        return (
            "A choropleth map displays country-level data on an interactive world map. "
            "Countries are shaded according to an aggregated value using a color scale."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        """
        Returns query parameters for the main grouped aggregation query.

        :return: A dictionary with a single ``"main"`` key containing the query parameters.
        :rtype: dict[str, QueryParams]
        """
        return {
            "main": QueryParams(
                group_fields=[self._country_variable],
                aggregations=[
                    TherismosAggregation(
                        id=self._aggregation.name,
                        function=AggregationFunction(self._aggregation.function),
                        field=self._aggregation.field,
                    ),
                ],
                dataframe_schema=pl.Schema(
                    self._get_schema_fields(),
                ),
            )
        }

    def _get_schema_fields(self) -> dict[str, type[pl.DataType]]:
        return {
            self._country_variable: pl.Utf8,
            self._aggregation.name: (pl.Int64 if self._aggregation.function == AggregationFunction.COUNT else pl.Float64),
        }

    def _prepare_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        df: pl.DataFrame = dfs["main"]
        df = df.sort(self._country_variable, descending=False)

        country_codes = df[self._country_variable].to_list()
        z_values = df[self._aggregation.name].to_list()
        country_names = [c.name if (c := pycountry.countries.get(alpha_2=code)) is not None else code for code in country_codes]
        alpha3_codes = [(c := pycountry.countries.get(alpha_2=code)) and c.alpha_3 or code for code in country_codes]

        use_discrete = bool(self._color_rules and any(not r.assigner.continuous for r in self._color_rules))

        discrete_colors: list[str] | None = None
        color_scale: HSLColorScale | None = None

        if use_discrete:
            current_colors: HSLColorSequence = [theme.get_color("primary_color_normal")] * len(country_codes)
            for color_rule in self._color_rules:
                if color_rule.assigner.continuous:
                    continue
                mask = color_rule.selector.selector(df, self._country_variable, self._aggregation.name)
                assigned = color_rule.assigner.assign_colors(df, self._country_variable, self._aggregation.name)
                current_colors = Palette.mix_color_sequences(current_colors, assigned, mask, indexing=Palette.MixIndexing.BY_POSITION)
            discrete_colors = [c.to_css_hex() for c in current_colors]
        else:
            color_sequence: HSLColorSequence
            if self._color_rules:
                color_sequence = self._color_rules[0].assigner.anchor_colors
            else:
                raw_sequence = theme.get_default_color_sequence()
                if len(raw_sequence) < 2:
                    color_sequence = [HSLColor(0, 0, 0.96, 1.0)] + list(raw_sequence)
                else:
                    color_sequence = raw_sequence
            color_scale = Palette.to_color_scale(color_sequence)

        return {
            "country_codes": country_codes,
            "country_names": country_names,
            "values": z_values,
            "_alpha3_codes": alpha3_codes,
            "_use_discrete": use_discrete,
            "_discrete_colors": discrete_colors,
            "_color_scale": color_scale,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        country_names = payload["country_names"]
        z_values = payload["values"]
        alpha3_codes = payload["_alpha3_codes"]
        use_discrete = payload["_use_discrete"]
        discrete_colors = payload["_discrete_colors"]
        color_scale = payload["_color_scale"]

        effective_title = payload["title"] or (f"Choropleth map of {self._aggregation.name} by {self._country_variable}")
        effective_value = self._value_label or self._append_unit(self._aggregation.default_label, self._aggregation.field)
        effective_country = self._country_label or self._country_variable
        hover = f"<b>%{{customdata}}</b><br>{effective_country}: %{{location}}<br>{effective_value}: %{{z}}<extra></extra>"

        if use_discrete:
            color_groups: dict[str, dict[str, Any]] = {}
            for a3, z_val, name, hex_val in zip(alpha3_codes, z_values, country_names, discrete_colors):
                group = color_groups.setdefault(hex_val, {"locations": [], "z": [], "names": []})
                group["locations"].append(a3)
                group["z"].append(z_val)
                group["names"].append(name)
            fig = go.Figure()
            for hex_val, group in color_groups.items():
                fig.add_trace(
                    go.Choropleth(
                        locationmode="ISO-3",
                        locations=group["locations"],
                        z=group["z"],
                        colorscale=[[0, hex_val], [1, hex_val]],
                        showscale=False,
                        customdata=group["names"],
                        hovertemplate=hover,
                    )
                )
        else:
            fig = go.Figure(
                data=go.Choropleth(
                    locationmode="ISO-3",
                    locations=alpha3_codes,
                    z=z_values,
                    colorscale=Palette.color_scale_to_plotly(color_scale),
                    customdata=country_names,
                    hovertemplate=hover,
                )
            )
            fig.update_traces(colorbar_title_text=effective_value)
            for color_rule in self._color_rules:
                if isinstance(color_rule.assigner, ColorAssignerByValue) and color_rule.assigner.range_min is not None:
                    fig.update_traces(zmin=color_rule.assigner.range_min)
                    break
            for color_rule in self._color_rules:
                if isinstance(color_rule.assigner, ColorAssignerByValue) and color_rule.assigner.range_max is not None:
                    fig.update_traces(zmax=color_rule.assigner.range_max)
                    break
            for color_rule in self._color_rules:
                if isinstance(color_rule.assigner, ColorAssignerByValue) and color_rule.assigner.range_center_at is not None:
                    fig.update_traces(zmid=color_rule.assigner.range_center_at)
                    break

        geo_settings: dict[str, Any] = dict(
            showframe=False,
            showcoastlines=True,
            projection_type=self._map_style.value,
            resolution=50 if self._resolution_high else 110,
            showcountries=self._show_countries,
            showlakes=self._show_countries,
            showocean=self._show_ocean,
        )
        if self._fit_bounds:
            geo_settings["fitbounds"] = "locations"
        fig.update_layout(autosize=True, height=600, geo=geo_settings)
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        return fig

    @classmethod
    def supported_color_selectors(cls) -> list[type[ColorSelectorStrategy]]:
        return [
            ColorSelectorStrategy,
            ColorSelectByCategoryName,
            ColorSelectValuesBelow,
            ColorSelectValuesAbove,
            ColorSelectMaxValue,
            ColorSelectMinValue,
            ColorSelectValuesBelowMean,
            ColorSelectValuesAboveMean,
        ]

    @classmethod
    def supported_color_assigners(cls) -> list[type[ColorAssignerStrategy]]:
        return [
            ColorAssignerByCategory,
            ColorAssignerByValue,
        ]


__all__ = [
    "ChoroplethMapChartStrategy",
]
