"""
Line chart figure strategy implementation.
"""

from itertools import groupby
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 plotly.subplots import make_subplots
from polychromos.palette import HSLColorSequence
from therismos.grouping import Aggregation as TherismosAggregation
from therismos.grouping import AggregationFunction

from orama.color import (
    ColorAssignerByCategory,
    ColorAssignerStrategy,
    ColorRule,
    ColorSelectByCategoryName,
    ColorSelectorStrategy,
)
from orama.enums import LineInterpolation, LinePlacement, LineSortOrder
from orama.options import (
    ChartOptionsMap,
    FloatOption,
    IntOption,
    LineInterpolationOption,
    LinePlacementOption,
    LineSortOrderOption,
)
from orama.query import QueryParams
from orama.strategies.base import FigureStrategy, _hex_to_rgb
from orama.theme import Theme
from orama.variables import (
    Aggregation,
    CategoricalVariable,
    ChartVariablesMap,
    DateTimeVariable,
    NumericalAggregationVariable,
)


class LineChartStrategy(FigureStrategy):
    """
    Line chart figure strategy.
    """

    _variables: ChartVariablesMap = {
        "date": DateTimeVariable(
            name="Date",
            description="Calendar variable. Documents will be aggregated by year and month.",
            map_to="date_variable",
        ),
        "category": CategoricalVariable(
            name="Category",
            description=("Categorical variable for the line chart. Each category corresponds to a line."),
            map_to="category_variable",
            multicategorical=False,
        ),
        "value": NumericalAggregationVariable(
            name="Value",
            description=("An aggregation of a numerical variable to produce the final value at each point in time."),
            map_to="aggregation",
        ),
    }
    _options: ChartOptionsMap = {
        "min_value": FloatOption(
            name="Minimum value",
            description="Minimum value for the line chart value axis.",
            map_to="min_value",
            default_value=None,
        ),
        "max_value": FloatOption(
            name="Maximum value",
            description="Maximum value for the line chart value axis.",
            map_to="max_value",
            default_value=None,
        ),
        "fill_opacity": FloatOption(
            name="Fill opacity",
            description="Opacity of the filled area under each line. 0 disables the fill.",
            map_to="fill_opacity",
            default_value=0.0,
            min_value=0.0,
            max_value=1.0,
        ),
        "line_width": IntOption(
            name="Line width",
            description="Width of each line in pixels.",
            map_to="line_width",
            default_value=3,
            min_value=1,
        ),
        "placement": LinePlacementOption(
            name="Placement",
            description="How the series lines/areas are positioned relative to each other.",
            map_to="placement",
            default_value=LinePlacement.OVERLAPPING,
        ),
        "interpolation": LineInterpolationOption(
            name="Interpolation",
            description="Line interpolation mode between data points.",
            map_to="interpolation",
            default_value=LineInterpolation.LINEAR,
        ),
        "sort_order": LineSortOrderOption(
            name="Sort series",
            description="Order in which the category lines are drawn.",
            map_to="sort_order",
            default_value=LineSortOrder.NONE,
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        date_variable: str,
        aggregation: Aggregation,
        category_variable: str | None = None,
        min_value: float | None = None,
        max_value: float | None = None,
        fill_opacity: float = 0.0,
        line_width: int = 3,
        placement: LinePlacement = LinePlacement.OVERLAPPING,
        interpolation: LineInterpolation = LineInterpolation.LINEAR,
        sort_order: LineSortOrder = LineSortOrder.NONE,
        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,
        date_label: str | None = None,
        value_label: str | None = None,
        category_label: str | None = None,
    ) -> None:
        super().__init__(
            fields_metadata=fields_metadata,
            color_rules=color_rules,
            title=title,
            subtitle=subtitle,
            caption=caption,
            description=description,
            storytelling_context=storytelling_context,
        )

        self._variables["date"].validate_fields(fields_metadata, date_variable)
        if aggregation.field:
            self._variables["value"].validate_fields(fields_metadata, aggregation.field)
        if category_variable:
            self._variables["category"].validate_fields(fields_metadata, category_variable)

        self._date_variable = date_variable
        self._aggregation = aggregation
        self._explicit_category_variable = category_variable
        self._category_variable = category_variable or "type"

        if min_value is not None:
            self._options["min_value"].validate(min_value)
        if max_value is not None:
            self._options["max_value"].validate(max_value)
        self._options["fill_opacity"].validate(fill_opacity)
        self._options["line_width"].validate(line_width)
        self._options["placement"].validate(placement)
        self._options["interpolation"].validate(interpolation)
        self._options["sort_order"].validate(sort_order)

        self._min_value = min_value
        self._max_value = max_value
        self._fill_opacity = fill_opacity
        self._line_width = line_width
        self._placement = placement
        self._interpolation = interpolation
        self._sort_order = sort_order
        self._date_label = date_label
        self._value_label = value_label
        self._category_label = category_label

    @classmethod
    def name(cls) -> str:
        return "Line Chart"

    @classmethod
    def description(cls) -> str:
        return (
            "A line chart represents information as a series of data points connected by straight "
            "line segments, used to visualize a trend in data over intervals of time."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        return {
            "main": QueryParams(
                group_fields=[
                    f"{self._date_variable}__year",
                    f"{self._date_variable}__month",
                    self._category_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,
        datetime_vars: list[str] | None = None,
        value_var: str | None = None,
    ) -> dict[str, type[pl.DataType]]:
        if not datetime_vars:
            datetime_vars = [
                f"{self._date_variable}__year",
                f"{self._date_variable}__month",
                self._category_variable,
            ]
        if not value_var:
            value_var = self._aggregation.name
        return {
            **{var: self._infer_column_type(self._fields_metadata, var) for var in datetime_vars},
            value_var: (pl.Int64 if self._aggregation.function == AggregationFunction.COUNT else pl.Float64),
        }

    def _build_line_dataframes(
        self,
        df: pl.DataFrame,
        tr: Translations,
    ) -> tuple[dict[str, pl.DataFrame], list[str], str]:
        """
        Build per-category DataFrames with date-range filling, null trimming, and date translation.

        :param df: Raw DataFrame from the query result.
        :type df: pl.DataFrame
        :param tr: Translations for localizing date labels.
        :type tr: Translations
        :return: Mapping of category to its processed DataFrame, ordered category list, and date alias.
        :rtype: tuple[dict[str, pl.DataFrame], list[str], str]
        """
        df = self._sort_and_translate(
            df,
            self._category_variable,
            self._fields_metadata[self._category_variable],
            tr,
            sort_asc=True,
        )
        line_dfs: dict[str, pl.DataFrame] = {}
        date_variables = (f"{self._date_variable}__year", f"{self._date_variable}__month")
        categories = [key for key, _ in groupby(df[self._category_variable].to_list())]
        alias = ""
        for category in categories:
            _df = df.filter(pl.col(self._category_variable) == category).select([*date_variables, self._aggregation.name])
            if not _df.height:
                continue
            for year in range(_df[date_variables[0]].min(), _df[date_variables[0]].max() + 1):  # type: ignore[arg-type, operator]
                for month in range(1, 13):
                    if not ((_df[date_variables[0]] == year) & (_df[date_variables[1]] == month)).any():
                        new_row = pl.DataFrame(
                            {
                                date_variables[0]: [year],
                                date_variables[1]: [month],
                                self._aggregation.name: [None],
                            }
                        )
                        _df = pl.concat([_df, new_row], how="vertical")
            _df = _df.sort(by=list(date_variables))
            _df = _df.with_row_index("__row_nr")
            agg_col = self._aggregation.name
            non_null = _df.filter(pl.col(agg_col).is_not_null())
            if non_null.height == 0:
                _df = _df.head(0)
            else:
                start = int(non_null["__row_nr"].min())  # type: ignore[arg-type]
                end = int(non_null["__row_nr"].max())  # type: ignore[arg-type]
                _df = _df.filter((pl.col("__row_nr") >= start) & (pl.col("__row_nr") <= end))
            if "__row_nr" in _df.columns:
                _df = _df.drop("__row_nr")
            for category_variable in reversed(date_variables):
                _df = self._sort_and_translate(
                    _df,
                    category_variable,
                    self._fields_metadata[category_variable],
                    tr,
                    sort_asc=True,
                )
            _df, alias = self._concatenate_variables(_df, list(date_variables))
            line_dfs[category] = _df
        return line_dfs, categories, alias

    def _compute_line_colors(
        self,
        categories_df: pl.DataFrame,
        theme: Theme,
    ) -> HSLColorSequence:
        """
        Produce a per-category color sequence via color rules or the default category assigner.

        :param categories_df: Single-column DataFrame of category values.
        :type categories_df: pl.DataFrame
        :param theme: Theme providing color sequences.
        :type theme: Theme
        :return: A color sequence with one entry per category.
        :rtype: HSLColorSequence
        """
        if self._color_rules:
            return self._apply_color_rules(categories_df, self._category_variable, self._aggregation.name, theme)
        return ColorAssignerByCategory(
            anchor_colors=theme.get_default_color_sequence(),
        ).assign_colors(categories_df, self._category_variable, self._aggregation.name)

    def _compute_value_range(
        self,
        df: pl.DataFrame,
    ) -> tuple[float, float]:
        """
        Derive the y-axis range from explicit overrides or the data extremes with a 5 % margin.

        :param df: The full (unfiltered) DataFrame.
        :type df: pl.DataFrame
        :return: ``(min_value, max_value)`` tuple for the y-axis range.
        :rtype: tuple[float, float]
        """
        min_val = (
            self._min_value
            if self._min_value is not None
            else (
                float(df[self._aggregation.name].min())  # type: ignore[arg-type]
                - 0.05 * abs(float(df[self._aggregation.name].cast(pl.Float64).min()))  # type: ignore[arg-type]
            )
        )
        max_val = (
            self._max_value
            if self._max_value is not None
            else (
                float(df[self._aggregation.name].max())  # type: ignore[arg-type]
                + 0.05 * abs(float(df[self._aggregation.name].cast(pl.Float64).max()))  # type: ignore[arg-type]
            )
        )
        return min_val, max_val

    def _sort_active_categories(
        self,
        active_categories: list[str],
        line_dfs: dict[str, pl.DataFrame],
    ) -> list[str]:
        """
        Re-order active categories according to :attr:`_sort_order`.

        :param active_categories: Categories present in ``line_dfs``.
        :type active_categories: list[str]
        :param line_dfs: Per-category DataFrames used to read first/last values.
        :type line_dfs: dict[str, pl.DataFrame]
        :return: Sorted copy of ``active_categories``.
        :rtype: list[str]
        """
        if self._sort_order == LineSortOrder.NONE:
            return active_categories
        if self._sort_order == LineSortOrder.ALPHABETICAL_ASC:
            return sorted(active_categories, key=str)
        if self._sort_order == LineSortOrder.ALPHABETICAL_DESC:
            return sorted(active_categories, key=str, reverse=True)
        if self._sort_order in (LineSortOrder.FIRST_VALUE_ASC, LineSortOrder.FIRST_VALUE_DESC):

            def _first(c: str) -> float:
                vals = line_dfs[c][self._aggregation.name].drop_nulls()
                return float(vals[0]) if vals.len() > 0 else 0.0

            return sorted(
                active_categories,
                key=_first,
                reverse=(self._sort_order == LineSortOrder.FIRST_VALUE_DESC),
            )
        if self._sort_order in (LineSortOrder.LAST_VALUE_ASC, LineSortOrder.LAST_VALUE_DESC):

            def _last(c: str) -> float:
                vals = line_dfs[c][self._aggregation.name].drop_nulls()
                return float(vals[-1]) if vals.len() > 0 else 0.0

            return sorted(
                active_categories,
                key=_last,
                reverse=(self._sort_order == LineSortOrder.LAST_VALUE_DESC),
            )
        return active_categories

    def _build_ridgeline_figure(
        self,
        active_categories: list[str],
        categories: list[str],
        line_dfs: dict[str, pl.DataFrame],
        line_colors: HSLColorSequence,
        alias: str,
        value_range: tuple[float, float],
    ) -> go.Figure:
        """
        Build a ridgeline (vertically stacked, shared x-axis) figure.

        :param active_categories: Ordered subset of categories to render.
        :type active_categories: list[str]
        :param categories: Full category list (used for color index lookup).
        :type categories: list[str]
        :param line_dfs: Per-category DataFrames.
        :type line_dfs: dict[str, pl.DataFrame]
        :param line_colors: Per-category color sequence.
        :type line_colors: HSLColorSequence
        :param alias: Date axis column name.
        :type alias: str
        :param value_range: Y-axis range applied when explicit bounds are set.
        :type value_range: tuple[float, float]
        :return: Ridgeline figure.
        :rtype: go.Figure
        """
        fig = make_subplots(
            rows=len(active_categories),
            cols=1,
            shared_xaxes=True,
            vertical_spacing=0.02,
        )
        for idx, category in enumerate(active_categories):
            i = categories.index(category)
            r, g, b = _hex_to_rgb(line_colors[i].to_css_hex())
            fill_kwargs: dict[str, Any] = {}
            if self._fill_opacity > 0:
                fill_kwargs = dict(
                    fill="tozeroy",
                    fillcolor=f"rgba({r},{g},{b},{self._fill_opacity})",
                )
            fig.add_trace(
                go.Scatter(
                    x=line_dfs[category][alias],
                    y=line_dfs[category][self._aggregation.name],
                    mode="lines",
                    line=dict(color=line_colors[i].to_css_hex(), width=self._line_width, shape=self._interpolation.value),
                    connectgaps=True,
                    name=str(category),
                    **fill_kwargs,
                ),
                row=idx + 1,
                col=1,
            )
            if self._min_value is not None or self._max_value is not None:
                fig.update_yaxes(range=value_range, row=idx + 1, col=1)
        return fig

    def _build_standard_figure(
        self,
        active_categories: list[str],
        categories: list[str],
        line_dfs: dict[str, pl.DataFrame],
        line_colors: HSLColorSequence,
        alias: str,
        value_range: tuple[float, float],
    ) -> go.Figure:
        """
        Build an overlapping, stacked, or stacked-relative figure.

        :param active_categories: Ordered subset of categories to render.
        :type active_categories: list[str]
        :param categories: Full category list (used for color index lookup).
        :type categories: list[str]
        :param line_dfs: Per-category DataFrames.
        :type line_dfs: dict[str, pl.DataFrame]
        :param line_colors: Per-category color sequence.
        :type line_colors: HSLColorSequence
        :param alias: Date axis column name.
        :type alias: str
        :param value_range: Y-axis range applied when applicable.
        :type value_range: tuple[float, float]
        :return: Standard line figure.
        :rtype: go.Figure
        """
        traces = []
        for category in active_categories:
            i = categories.index(category)
            r, g, b = _hex_to_rgb(line_colors[i].to_css_hex())
            if self._placement == LinePlacement.OVERLAPPING:
                extra: dict[str, Any] = {}
                if self._fill_opacity > 0:
                    extra = dict(
                        fill="tozeroy",
                        fillcolor=f"rgba({r},{g},{b},{self._fill_opacity})",
                    )
            elif self._placement == LinePlacement.STACKED:
                extra = dict(
                    stackgroup="one",
                    fillcolor=f"rgba({r},{g},{b},{self._fill_opacity})",
                )
            else:
                extra = dict(
                    stackgroup="one",
                    groupnorm="percent",
                    fillcolor=f"rgba({r},{g},{b},{self._fill_opacity})",
                )
            y_data = line_dfs[category][self._aggregation.name]
            if self._placement != LinePlacement.OVERLAPPING:
                mask = y_data.is_not_null()
                x_data = line_dfs[category][alias].filter(mask)
                y_data = y_data.filter(mask)
            else:
                x_data = line_dfs[category][alias]
            traces.append(
                go.Scatter(
                    x=x_data,
                    y=y_data,
                    mode="lines",
                    line=dict(color=line_colors[i].to_css_hex(), width=self._line_width, shape=self._interpolation.value),
                    connectgaps=True,
                    name=str(category),
                    **extra,
                )
            )
        fig = go.Figure(data=traces)
        if self._placement == LinePlacement.OVERLAPPING:
            fig.update_yaxes(range=value_range)
        elif self._min_value is not None or self._max_value is not None:
            fig.update_yaxes(range=value_range)
        return fig

    def _configure_line_layout(
        self,
        fig: go.Figure,
        active_categories: list[str],
        categories: list[str],
        line_colors: HSLColorSequence,
        alias: str,
        effective_date: str,
        effective_value: str,
    ) -> None:
        """
        Apply axis titles, hover mode, legend, and ridgeline annotations.

        :param fig: Plotly figure to update.
        :type fig: go.Figure
        :param active_categories: Ordered subset of categories rendered.
        :type active_categories: list[str]
        :param categories: Full category list (used for color index lookup).
        :type categories: list[str]
        :param line_colors: Per-category color sequence.
        :type line_colors: HSLColorSequence
        :param alias: Date axis column name (unused for standard layout but kept for consistency).
        :type alias: str
        :param effective_date: Resolved date axis label.
        :type effective_date: str
        :param effective_value: Resolved value axis label.
        :type effective_value: str
        """
        if self._placement == LinePlacement.RIDGELINE:
            n_rows = len(active_categories)
            fig.update_xaxes(title_text=effective_date, row=n_rows, col=1)
            fig.update_xaxes(showgrid=False)
            fig.update_yaxes(showgrid=False, showticklabels=False)
            fig.update_layout(hovermode="x", showlegend=False)
            for idx, category in enumerate(active_categories):
                i = categories.index(category)
                axis_key = "yaxis" if idx == 0 else f"yaxis{idx + 1}"
                domain = fig.layout[axis_key].domain
                y_center = (domain[0] + domain[1]) / 2
                fig.add_annotation(
                    text=str(category),
                    xref="paper",
                    yref="paper",
                    x=1,
                    y=y_center,
                    xanchor="left",
                    yanchor="middle",
                    showarrow=False,
                    font=dict(color=line_colors[i].to_css_hex()),
                )
            top_domain = fig.layout["yaxis"].domain[1]
            last_axis_key = "yaxis" if n_rows == 1 else f"yaxis{n_rows}"
            bottom_domain = fig.layout[last_axis_key].domain[0]
            y_all_center = (top_domain + bottom_domain) / 2
            fig.add_annotation(
                text=effective_value,
                xref="paper",
                yref="paper",
                x=-0.07,
                y=y_all_center,
                xanchor="center",
                yanchor="middle",
                textangle=-90,
                showarrow=False,
            )
        else:
            fig.update_layout(
                xaxis_title=effective_date,
                yaxis_title=effective_value,
                hovermode="x",
                xaxis=dict(showgrid=False),
            )
            if self._explicit_category_variable is not None:
                fig.update_layout(legend_title_text=self._category_label or self._category_variable)

    def _prepare_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        df = dfs["main"]
        line_dfs, categories, alias = self._build_line_dataframes(df, tr)
        categories_df = pl.DataFrame({self._category_variable: categories})
        line_colors = self._compute_line_colors(categories_df, theme)
        value_range = self._compute_value_range(df)
        active_categories = self._sort_active_categories([c for c in categories if c in line_dfs], line_dfs)
        return {
            "categories": active_categories,
            "colors": [line_colors[categories.index(c)].to_css_hex() for c in active_categories],
            "_line_dfs": line_dfs,
            "_all_categories": categories,
            "_line_colors": line_colors,
            "_alias": alias,
            "_value_range": value_range,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        active_categories = payload["categories"]
        categories = payload["_all_categories"]
        line_dfs = payload["_line_dfs"]
        line_colors = payload["_line_colors"]
        alias = payload["_alias"]
        value_range = payload["_value_range"]
        if self._placement == LinePlacement.RIDGELINE:
            fig = self._build_ridgeline_figure(active_categories, categories, line_dfs, line_colors, alias, value_range)
        else:
            fig = self._build_standard_figure(active_categories, categories, line_dfs, line_colors, alias, value_range)
        effective_title = payload["title"] or (f"Line chart of {self._aggregation.name} by {self._date_variable}")
        effective_date = self._date_label or self._date_variable
        effective_value = self._value_label or self._append_unit(self._aggregation.default_label, self._aggregation.field)
        self._configure_line_layout(fig, active_categories, categories, line_colors, alias, effective_date, effective_value)
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        return fig

    @classmethod
    def supported_color_selectors(cls) -> list[type[ColorSelectorStrategy]]:
        return [
            ColorSelectorStrategy,
            ColorSelectByCategoryName,
        ]

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


__all__ = [
    "LineChartStrategy",
]
