"""Utilitaires pour générer des heatmaps calendrier configurables."""

from __future__ import annotations

import calendar
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Sequence

import numpy as np
import pandas as pd

from meteo.analysis import compute_daily_rainfall_totals

from .calendar import plot_calendar_heatmap

CalendarAggregator = Callable[[pd.DataFrame], pd.Series]

__all__ = [
    "CalendarHeatmapSpec",
    "CalendarHeatmapResult",
    "daily_mean_series",
    "rainfall_daily_total_series",
    "format_calendar_matrix",
    "generate_calendar_heatmaps",
]


@dataclass(frozen=True)
class CalendarHeatmapSpec:
    """Description d'une heatmap calendrier à générer."""

    key: str
    agg_label: str
    title_template: str
    cmap: str
    colorbar_label: str
    aggregator: CalendarAggregator


@dataclass(frozen=True)
class CalendarHeatmapResult:
    """Résultat d'une tentative de génération de heatmap calendrier."""

    spec: CalendarHeatmapSpec
    output_path: Path | None
    reason: str | None = None


def daily_mean_series(column: str) -> CalendarAggregator:
    """Retourne un agrégateur qui calcule la moyenne quotidienne d'une colonne."""

    def _aggregator(df: pd.DataFrame) -> pd.Series:
        if column not in df.columns:
            raise KeyError(column)
        if not isinstance(df.index, pd.DatetimeIndex):
            raise TypeError("Le DataFrame doit être indexé par des timestamps pour les moyennes quotidiennes.")
        return df[column].resample("1D").mean()

    return _aggregator


def rainfall_daily_total_series(df: pd.DataFrame) -> pd.Series:
    """Calcule les cumuls quotidiens de précipitations."""

    totals = compute_daily_rainfall_totals(df=df)
    return totals["daily_total"]


def format_calendar_matrix(series: pd.Series, year: int, agg_label: str) -> pd.DataFrame:
    """Transforme une série quotidienne en matrice calendrier mois x jours."""

    if not isinstance(series.index, pd.DatetimeIndex):
        raise TypeError("La série doit avoir un index temporel pour être convertie en calendrier.")

    tz = series.index.tz
    start = pd.Timestamp(year=year, month=1, day=1, tz=tz)
    end = pd.Timestamp(year=year, month=12, day=31, tz=tz)
    filtered = series.loc[(series.index >= start) & (series.index <= end)]

    matrix = pd.DataFrame(
        np.nan,
        index=[calendar.month_name[m][:3] for m in range(1, 13)],
        columns=list(range(1, 32)),
    )

    for timestamp, value in filtered.items():
        matrix.at[calendar.month_name[timestamp.month][:3], timestamp.day] = value

    matrix.index.name = f"{agg_label} ({year})"
    return matrix


def generate_calendar_heatmaps(
    df: pd.DataFrame,
    specs: Sequence[CalendarHeatmapSpec],
    *,
    year: int,
    output_dir: str | Path,
) -> list[CalendarHeatmapResult]:
    """Génère l'ensemble des heatmaps calendrier décrites dans `specs`."""

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    results: list[CalendarHeatmapResult] = []
    for spec in specs:
        try:
            daily_series = spec.aggregator(df)
        except Exception as exc:
            results.append(
                CalendarHeatmapResult(
                    spec=spec,
                    output_path=None,
                    reason=str(exc),
                )
            )
            continue

        if daily_series is None or daily_series.empty:
            results.append(
                CalendarHeatmapResult(
                    spec=spec,
                    output_path=None,
                    reason="Série vide ou invalide.",
                )
            )
            continue

        try:
            matrix = format_calendar_matrix(daily_series, year, spec.agg_label)
        except Exception as exc:  # pragma: no cover - remonté au résultat
            results.append(
                CalendarHeatmapResult(
                    spec=spec,
                    output_path=None,
                    reason=str(exc),
                )
            )
            continue

        output_path = output_dir / f"calendar_{spec.key}_{year}.png"
        plot_calendar_heatmap(
            matrix=matrix,
            output_path=output_path,
            title=spec.title_template.format(year=year, label=spec.agg_label),
            cmap=spec.cmap,
            colorbar_label=spec.colorbar_label,
        )
        results.append(CalendarHeatmapResult(spec=spec, output_path=output_path.resolve()))

    return results