"""
Abstract base class for figure strategies.
"""

import datetime
from abc import ABC, abstractmethod
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

import plotly.graph_objects as go
import polars as pl
from babel.support import Translations
from fields_metadata import FieldMetadata
from polychromos.palette import HSLColorSequence, Palette

from orama.color import (
    ColorAssignerStrategy,
    ColorRule,
    ColorSelectorStrategy,
    ColorStrategyDescription,
)
from orama.options import ChartOptionsMap
from orama.query import QueryParams
from orama.theme import Theme
from orama.variables import ChartVariablesMap


def _hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
    hex_color = hex_color.lstrip("#")
    return int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)


_PYTHON_TYPE_TO_POLARS: dict[type, type[pl.DataType]] = {
    int: pl.Int64,
    float: pl.Float64,
    str: pl.Utf8,
    bool: pl.Boolean,
    datetime.datetime: pl.Datetime,
    datetime.date: pl.Datetime,
}


@dataclass(frozen=True)
class FigureStrategyDescription:
    """
    Describes a figure strategy.
    """

    name: str
    description: str
    fields_metadata: dict[str, FieldMetadata]
    colors: ColorStrategyDescription
    variables: dict[str, dict[str, Any]] = field(default_factory=dict)
    options: dict[str, dict[str, Any]] = field(default_factory=dict)
    icon: str | None = None
    max_color_rules: int | None = None


@dataclass
class FigureView:
    """
    The result of a figure strategy's ``plot()`` call.

    Bundles the rendered Plotly figure together with its optional caption and
    description so callers receive everything they need to embed the chart in
    one object.

    :param figure: The rendered Plotly figure.
    :type figure: go.Figure
    :param caption: Optional caption text, intended to be rendered as a
        ``<figcaption>`` element when embedding in an HTML ``<figure>`` wrapper.
    :type caption: str | None
    :param description: Optional free-form text explaining how to interpret the
        figure (e.g., methodology notes, caveats). Not applied to the Plotly
        figure itself — returned for the caller to render as desired.
    :type description: str | None
    """

    figure: go.Figure
    caption: str | None
    description: str | None
    storytelling_context: str | None


class FigureStrategy(ABC):
    """
    Abstract base class for figure strategies.
    """

    _variables: ChartVariablesMap = {}
    _options: ChartOptionsMap = {}

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        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,
        pre_build_callback: Callable[[dict[str, Any]], None] | None = None,
        post_build_callback: Callable[[dict[str, Any]], None] | None = None,
    ) -> None:
        for color_rule in color_rules or []:
            if color_rule.selector.__class__ not in self.supported_color_selectors():
                raise ValueError(
                    f"Color selector {color_rule.selector.__class__.__name__} is not supported by {self.__class__.__name__}"
                )
            if color_rule.assigner.__class__ not in self.supported_color_assigners():
                raise ValueError(
                    f"Color assigner {color_rule.assigner.__class__.__name__} is not supported by {self.__class__.__name__}"
                )

        self._fields_metadata = fields_metadata
        self._color_rules = color_rules if color_rules is not None else []
        self._title = title
        self._subtitle = subtitle
        self._caption = caption
        self._description = description
        self._storytelling_context = storytelling_context
        self._pre_build_callback = pre_build_callback
        self._post_build_callback = post_build_callback

    @classmethod
    def name(cls) -> str:
        """
        The name of the figure strategy.

        :return: The name of the strategy.
        :rtype: str
        """
        return cls.__name__

    @classmethod
    def description(cls) -> str:
        """
        The description of the figure strategy.

        :return: The description of the strategy.
        :rtype: str
        """
        return ""

    @classmethod
    def icon(cls) -> str | None:
        """
        An optional SVG icon for the figure strategy.

        :return: An SVG string, or ``None`` if not set.
        :rtype: str | None
        """
        return None

    @classmethod
    def describe(
        cls,
        fields_metadata: dict[str, FieldMetadata],
    ) -> FigureStrategyDescription:
        """
        Describes the figure strategy with its variables, options, accepted coloring rules, and
        other metadata.

        :param fields_metadata: Mapping of field paths to their metadata.
        :type fields_metadata: dict[str, FieldMetadata]
        :return: A description of the figure strategy.
        :rtype: FigureStrategyDescription
        """
        return FigureStrategyDescription(
            name=cls.name(),
            description=cls.description(),
            icon=cls.icon(),
            fields_metadata=fields_metadata,
            variables={var_name: var.describe(fields_metadata) for var_name, var in cls._variables.items()},
            options={opt_name: opt.describe() for opt_name, opt in cls._options.items()},
            colors=ColorStrategyDescription(
                supported_selectors=list(cls.supported_color_selectors()),
                supported_assigners=list(cls.supported_color_assigners()),
            ),
            max_color_rules=cls.max_color_rules(),
        )

    @abstractmethod
    def get_query_params(self) -> dict[str, QueryParams]:
        """
        Gets the parameters to use for grouped aggregation queries.

        :return: The query parameters keyed by query name.
        :rtype: dict[str, QueryParams]
        """

    def _prepare_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        """
        Prepare strategy-specific data for rendering.

        Called by :meth:`build` before the ``pre_build_callback`` and by :meth:`plot`
        before rendering. The returned dict is merged into the payload that the callback
        receives. Keys without a leading underscore are the public, documented payload
        keys (e.g. ``categories``, ``colors``); keys with a leading underscore (e.g.
        ``_df``, ``_alias``) are internal rendering state consumed by
        :meth:`_render_from_payload`.

        Callback mutations to public keys **do** affect the rendered figure, because
        :meth:`_render_from_payload` reads directly from the (possibly mutated) payload.

        Override in each strategy. Default returns an empty dict.

        :param dfs: The DataFrames containing the data to render.
        :type dfs: dict[str, pl.DataFrame]
        :param tr: Translations for localizing text.
        :type tr: Translations
        :param theme: Theme for computing colors.
        :type theme: Theme
        :return: Strategy-specific prepared data to merge into the payload.
        :rtype: dict[str, Any]
        """
        return {}

    @abstractmethod
    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        """
        Build and return the Plotly figure from the prepared payload.

        Called by :meth:`build` after the ``pre_build_callback`` (which may have
        mutated the payload) and by :meth:`plot` directly. Should not call
        ``_apply_template`` — that is handled by the caller.

        :param payload: The merged payload produced by :meth:`_build_pre_payload`
            and :meth:`_prepare_plot_data`, possibly mutated by the callback.
        :type payload: dict[str, Any]
        :return: The rendered Plotly figure (without template applied).
        :rtype: go.Figure
        """

    def plot(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> FigureView:
        """
        Generates the plotly figure based on the data provided and the strategy specifics.

        :param dfs: The DataFrames containing the data to plot.
        :type dfs: dict[str, pl.DataFrame]
        :param tr: The translations to use for localizing text.
        :type tr: Translations
        :param theme: The theme to use for styling the figure.
        :type theme: Theme
        :return: A view containing the rendered figure and its optional caption.
        :rtype: FigureView
        """
        payload = self._build_pre_payload(dfs)
        payload.update(self._prepare_plot_data(dfs, tr, theme))
        fig = self._render_from_payload(payload)
        self._apply_template(fig, theme)
        return FigureView(
            figure=fig,
            caption=payload["caption"],
            description=payload["description"],
            storytelling_context=payload["storytelling_context"],
        )

    @property
    def title(self) -> str | None:
        """
        The figure title, if any.

        :return: The title string, or ``None`` if not set.
        :rtype: str | None
        """
        return self._title

    @property
    def subtitle(self) -> str | None:
        """
        The figure subtitle, if any.

        :return: The subtitle string, or ``None`` if not set.
        :rtype: str | None
        """
        return self._subtitle

    @property
    def caption(self) -> str | None:
        """
        The figure caption, if any. Intended to be rendered as a ``<figcaption>`` element
        by the caller when embedding the figure in an HTML ``<figure>`` wrapper.

        :return: The caption string, or ``None`` if not set.
        :rtype: str | None
        """
        return self._caption

    @property
    def storytelling_context(self) -> str | None:
        """
        Optional narrative context for the figure, intended to help build a story
        around the data.

        :return: The storytelling context string, or ``None`` if not set.
        :rtype: str | None
        """
        return self._storytelling_context

    def _append_unit(self, label: str, field: str | None) -> str:
        """
        Appends the unit from field metadata to a label, if defined.

        :param label: The base label to append the unit to.
        :type label: str
        :param field: The field name to look up in ``_fields_metadata``.
        :type field: str | None
        :return: The label with the unit appended as ``(unit)``, or the original label if no unit is set.
        :rtype: str
        """
        if field is None:
            return label
        metadata = self._fields_metadata.get(field)
        if metadata is None:
            return label
        unit = metadata.extra.get("unit") if metadata.extra else None
        if not unit:
            return label
        return f"{label} ({unit})"

    def _apply_title(
        self,
        fig: go.Figure,
        title: str,
        subtitle: str | None = None,
    ) -> None:
        if subtitle:
            fig.update_layout(title={"text": title, "subtitle": {"text": subtitle}})
        else:
            fig.update_layout(title=title)

    def _apply_template(
        self,
        fig: go.Figure,
        theme: Theme,
    ) -> None:
        fig.update_layout(
            template=theme.plotly_template(),
        )

    def _infer_column_type(
        self,
        fields_metadata: dict[str, FieldMetadata],
        column_field: str,
    ) -> type[pl.DataType]:
        if column_field not in fields_metadata:
            raise ValueError(f"Field {column_field} not found in fields metadata")
        field_type = fields_metadata[column_field].effective_type
        if field_type is None:
            return pl.Utf8
        return _PYTHON_TYPE_TO_POLARS.get(field_type, pl.Utf8)

    def _concatenate_variables(
        self,
        df: pl.DataFrame,
        variables: list[str],
        alias: str | None = None,
    ) -> tuple[pl.DataFrame, str]:
        if not variables:
            raise ValueError("No variables to concatenate")
        if not alias:
            alias = " - ".join(variables)
        out_df: pl.DataFrame
        if len(variables) == 1:
            out_df = df.rename({variables[0]: alias})
        else:
            out_df = df.with_columns(
                pl.concat_str(
                    variables,
                    separator=" - ",
                )
                .cast(pl.Utf8)
                .alias(alias)
            ).drop(variables)
        return out_df, alias

    def _get_temporal_label_mapping(
        self,
        field_metadata: FieldMetadata,
        tr: Translations,
    ) -> dict[int, str] | None:
        """
        Return a numeric-to-label mapping for month or day-of-week fields, or ``None``.

        :param field_metadata: Metadata for the field being translated.
        :type field_metadata: FieldMetadata
        :param tr: Translations used for localizing month/day names.
        :type tr: Translations
        :return: Mapping dict for month or dow fields, ``None`` otherwise.
        :rtype: dict[int, str] | None
        """
        validation = field_metadata.extra.get("suggested_validation") if field_metadata.extra else None
        if validation == "month":
            return {
                1: tr.gettext("January"),
                2: tr.gettext("February"),
                3: tr.gettext("March"),
                4: tr.gettext("April"),
                5: tr.gettext("May"),
                6: tr.gettext("June"),
                7: tr.gettext("July"),
                8: tr.gettext("August"),
                9: tr.gettext("September"),
                10: tr.gettext("October"),
                11: tr.gettext("November"),
                12: tr.gettext("December"),
            }
        if validation == "dow":
            return {
                1: tr.gettext("Monday"),
                2: tr.gettext("Tuesday"),
                3: tr.gettext("Wednesday"),
                4: tr.gettext("Thursday"),
                5: tr.gettext("Friday"),
                6: tr.gettext("Saturday"),
                7: tr.gettext("Sunday"),
            }
        return None

    def _sort_and_translate(
        self,
        df: pl.DataFrame,
        table_column: str,
        field_metadata: FieldMetadata,
        tr: Translations,
        sort_asc: bool | None = None,
    ) -> pl.DataFrame:
        if sort_asc is not None:
            df = df.sort(table_column, descending=not sort_asc)
        label_mapping = self._get_temporal_label_mapping(field_metadata, tr)
        validation = field_metadata.extra.get("suggested_validation") if field_metadata.extra else None
        if label_mapping is not None:
            df = df.with_columns(pl.col(table_column).cast(pl.Utf8).replace(label_mapping).alias(table_column))
        elif validation == "year":
            df = df.with_columns(pl.col(table_column).cast(pl.Utf8).alias(table_column))
        return df

    def _apply_color_rules(
        self,
        df: pl.DataFrame,
        category_column: str,
        value_column: str,
        theme: Theme,
    ) -> HSLColorSequence:
        """
        Apply all color rules to produce a per-row color sequence.

        :param df: DataFrame used by the color selectors and assigners.
        :type df: pl.DataFrame
        :param category_column: Column name used as the category axis.
        :type category_column: str
        :param value_column: Column name used as the value axis.
        :type value_column: str
        :param theme: Theme providing the base color.
        :type theme: Theme
        :return: A color sequence with one entry per row of ``df``.
        :rtype: HSLColorSequence
        """
        colors: HSLColorSequence = [theme.get_color("primary_color_normal")] * df.height
        for rule in self._color_rules:
            selector = rule.selector.selector(df, category_column, value_column)
            assigned = rule.assigner.assign_colors(df, category_column, value_column)
            colors = Palette.mix_color_sequences(colors, assigned, selector, indexing=Palette.MixIndexing.BY_POSITION)
        return colors

    @abstractmethod
    def _get_schema_fields(self) -> dict[str, type[pl.DataType]]:
        """
        Returns the Polars schema fields expected in the result DataFrame.

        Implementors must return a mapping of column name to Polars DataType class
        that reflects all fields produced by :meth:`get_query_params`. This schema
        is used to enforce type safety when constructing query parameters.

        :return: Mapping of column name to Polars type class.
        :rtype: dict[str, type[pl.DataType]]
        """

    @classmethod
    def max_color_rules(cls) -> int | None:
        """
        The maximum number of color rules supported by this strategy.

        :return: The maximum number of color rules, or ``None`` if unlimited.
        :rtype: int | None
        """
        return None

    @classmethod
    def supported_color_selectors(cls) -> list[type[ColorSelectorStrategy]]:
        """
        Returns the list of supported color selector strategies.

        :return: Supported selector classes; empty by default.
        :rtype: list[type[ColorSelectorStrategy]]
        """
        return []

    @classmethod
    def supported_color_assigners(cls) -> list[type[ColorAssignerStrategy]]:
        """
        Returns the list of supported color assigner strategies.

        :return: Supported assigner classes; empty by default.
        :rtype: list[type[ColorAssignerStrategy]]
        """
        return []

    def _serialise_color_rules(self) -> list[dict[str, Any]]:
        return [
            {
                "selector_type": rule.selector.type_id,
                "assigner_type": rule.assigner.type_id,
                "colors": [c.to_css_hex() for c in rule.assigner.anchor_colors],
            }
            for rule in self._color_rules
        ]

    def _build_base_metadata(self) -> dict[str, Any]:
        return {
            "strategy": type(self).__name__,
            "title": self._title,
            "subtitle": self._subtitle,
            "caption": self._caption,
            "description": self._description,
            "storytelling_context": self._storytelling_context,
            "color_rules": self._serialise_color_rules(),
        }

    def _build_pre_payload(self, dfs: dict[str, pl.DataFrame]) -> dict[str, Any]:
        return {
            "dfs": dict(dfs),
            "strategy": type(self).__name__,
            "title": self._title,
            "subtitle": self._subtitle,
            "caption": self._caption,
            "description": self._description,
            "storytelling_context": self._storytelling_context,
            "color_rules": [rule.explain() for rule in self._color_rules],
            "variables": {
                key: {
                    "name": var._name,
                    "description": var._description,
                    "value": getattr(self, f"_{var._map_to}", None),
                }
                for key, var in self._variables.items()
            },
            "options": {
                key: {
                    "name": opt._name,
                    "description": opt._description,
                    "value": getattr(self, f"_{opt._map_to}", None),
                }
                for key, opt in self._options.items()
            },
        }

    def _build_post_payload(self, view: FigureView, dfs: dict[str, pl.DataFrame]) -> dict[str, Any]:
        return {
            "figure": view.figure,
            "caption": view.caption,
            "description": view.description,
            "storytelling_context": view.storytelling_context,
            "metadata": self._build_base_metadata(),
            "data": {key: {"columns": df.columns, "rows": df.to_dicts()} for key, df in dfs.items()},
        }

    def build(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> FigureView:
        """
        Builds the figure, honouring any registered pre/post-build callbacks.

        This is the recommended entry point for rendering a figure. It delegates
        to :meth:`plot` for the actual rendering and wraps it with optional
        callback hooks:

        - ``pre_build_callback`` fires *before* :meth:`plot` and receives a flat
          dict with the aggregated DataFrames and strategy metadata. Mutations to
          ``payload["dfs"]`` change what gets plotted. Mutations to
          ``payload["title"]``, ``payload["subtitle"]``, ``payload["caption"]``,
          ``payload["description"]``, and ``payload["storytelling_context"]`` are
          reflected in the rendered figure and returned :class:`FigureView`.
        - ``post_build_callback`` fires *after* :meth:`plot` and receives the
          rendered figure, caption, description, and data snapshot. Mutations
          to ``payload["figure"]``, ``payload["caption"]``, or
          ``payload["description"]`` are reflected in the returned
          :class:`FigureView`.

        When no callbacks are set, ``build()`` is equivalent to ``plot()``.

        :param dfs: Aggregated DataFrames keyed by query name.
        :type dfs: dict[str, pl.DataFrame]
        :param tr: Translations for i18n text.
        :type tr: Translations
        :param theme: Theme for styling the figure.
        :type theme: Theme
        :return: The rendered figure view.
        :rtype: FigureView
        """
        payload = self._build_pre_payload(dfs)
        payload.update(self._prepare_plot_data(payload["dfs"], tr, theme))
        if self._pre_build_callback is not None:
            self._pre_build_callback(payload)
        fig = self._render_from_payload(payload)
        self._apply_template(fig, theme)
        view = FigureView(
            figure=fig,
            caption=payload["caption"],
            description=payload["description"],
            storytelling_context=payload["storytelling_context"],
        )

        if self._post_build_callback is None:
            return view

        post = self._build_post_payload(view, payload["dfs"])
        self._post_build_callback(post)
        return FigureView(
            figure=post["figure"],
            caption=post["caption"],
            description=post["description"],
            storytelling_context=post["storytelling_context"],
        )

    def modebar_config(self) -> dict[str, Any]:
        """
        Returns the serialisable Plotly config dict applied to each figure.

        The returned dict is passed directly to ``plotly.Figure.to_html(config=...)``.
        It covers settings that do not require JavaScript callbacks — ``displaylogo``
        and ``modeBarButtonsToRemove``.

        Override in a subclass to customise the modebar for a specific strategy.

        :return: Plotly config dict.
        :rtype: dict[str, Any]
        """
        return {
            "displaylogo": False,
            "modeBarButtonsToRemove": [
                "select2d",
                "lasso2d",
            ],
        }


__all__ = [
    "FigureStrategy",
    "FigureStrategyDescription",
]
