"""Fonctions utilitaires pour générer des hexbins à partir de scénarios."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Sequence

import pandas as pd

from meteo.correlation_presets import HexbinScenario
from meteo.variables import VARIABLES_BY_KEY, Variable

from .relationships import plot_hexbin_with_third_variable

__all__ = [
    "HexbinPlotResult",
    "prepare_hexbin_dataframe",
    "generate_hexbin_for_scenario",
    "generate_hexbin_scenarios",
]


@dataclass(frozen=True)
class HexbinPlotResult:
    """Résultat d'un hexbin généré à partir d'un scénario prédéfini."""

    scenario: HexbinScenario
    var_x: Variable
    var_y: Variable
    var_color: Variable
    output_path: Path
    point_count: int


def _get_variables_for_scenario(scenario: HexbinScenario) -> tuple[Variable, Variable, Variable]:
    try:
        return (
            VARIABLES_BY_KEY[scenario.key_x],
            VARIABLES_BY_KEY[scenario.key_y],
            VARIABLES_BY_KEY[scenario.key_color],
        )
    except KeyError as exc:
        raise KeyError(f"Variable inconnue dans le scénario hexbin : {exc}") from exc


def prepare_hexbin_dataframe(
    df: pd.DataFrame,
    var_x: Variable,
    var_y: Variable,
    var_color: Variable,
) -> pd.DataFrame:
    """Sélectionne et nettoie les colonnes utiles à un hexbin."""

    columns = [var_x.column, var_y.column, var_color.column]
    missing = [col for col in columns if col not in df.columns]
    if missing:
        raise KeyError(f"Colonnes absentes dans le DataFrame : {missing}")

    return df[columns].dropna()


def generate_hexbin_for_scenario(
    df: pd.DataFrame,
    scenario: HexbinScenario,
    *,
    base_output_dir: str | Path,
    cmap: str = "viridis",
) -> HexbinPlotResult:
    """Construit un hexbin complet à partir d'un scénario et retourne ses métadonnées."""

    var_x, var_y, var_color = _get_variables_for_scenario(scenario)
    cleaned_df = prepare_hexbin_dataframe(df, var_x, var_y, var_color)

    output_dir = Path(base_output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / scenario.filename

    reduce_func = scenario.get_reduce_func()
    reduce_label = scenario.get_reduce_label()

    plot_hexbin_with_third_variable(
        df=cleaned_df,
        var_x=var_x,
        var_y=var_y,
        var_color=var_color,
        output_path=output_path,
        gridsize=scenario.gridsize,
        mincnt=scenario.mincnt,
        reduce_func=reduce_func,
        reduce_func_label=reduce_label,
        cmap=cmap,
    )

    return HexbinPlotResult(
        scenario=scenario,
        var_x=var_x,
        var_y=var_y,
        var_color=var_color,
        output_path=output_path.resolve(),
        point_count=len(cleaned_df),
    )


def generate_hexbin_scenarios(
    df: pd.DataFrame,
    scenarios: Sequence[HexbinScenario],
    *,
    base_output_dir: str | Path,
    cmap: str = "viridis",
) -> list[HexbinPlotResult]:
    """Génère en série plusieurs hexbins prédéfinis et retourne les résultats."""

    if not scenarios:
        raise ValueError("Aucun scénario hexbin fourni.")

    results: list[HexbinPlotResult] = []
    for scenario in scenarios:
        result = generate_hexbin_for_scenario(
            df=df,
            scenario=scenario,
            base_output_dir=base_output_dir,
            cmap=cmap,
        )
        results.append(result)

    return results