# meteo/solar.py
from __future__ import annotations

import pandas as pd
from astral import Observer
from astral.sun import elevation


def _ensure_datetime_index(index: pd.Index) -> pd.DatetimeIndex:
    if not isinstance(index, pd.DatetimeIndex):
        raise TypeError("Un DatetimeIndex est requis pour calculer l'élévation solaire.")
    return index


def _prepare_index(index: pd.DatetimeIndex) -> pd.DatetimeIndex:
    """
    Retourne une version timezone-aware (en UTC) du DatetimeIndex fourni.
    """
    if index.tz is None:
        return index.tz_localize("UTC")
    return index.tz_convert("UTC")


def compute_solar_elevation_series(
    index: pd.Index,
    *,
    latitude: float,
    longitude: float,
    elevation_m: float = 0.0,
    series_name: str = "sun_elevation",
) -> pd.Series:
    """
    Calcule l'élévation du soleil (en degrés) pour chaque timestamp de l'index.
    """
    dt_index = _ensure_datetime_index(index)
    observer = Observer(latitude=latitude, longitude=longitude, elevation=elevation_m)
    utc_index = _prepare_index(dt_index)

    values = [
        float(elevation(observer, ts.to_pydatetime()))
        for ts in utc_index
    ]

    return pd.Series(values, index=dt_index, name=series_name)


def add_solar_elevation_column(
    df: pd.DataFrame,
    *,
    latitude: float,
    longitude: float,
    elevation_m: float = 0.0,
    column_name: str = "sun_elevation",
) -> pd.DataFrame:
    """
    Ajoute une colonne `column_name` contenant l'élévation du soleil en degrés.
    """
    series = compute_solar_elevation_series(
        df.index,
        latitude=latitude,
        longitude=longitude,
        elevation_m=elevation_m,
        series_name=column_name,
    )
    df[column_name] = series
    return df