"""Meteocat client."""

from collections.abc import Mapping

import pandas as pd
from pyregeon import RegionType

from meteora import settings, utils
from meteora.clients.base import BaseJSONClient
from meteora.clients.mixins import (
    APIKeyHeaderMixin,
    StationsEndpointMixin,
    TimePartitionedTSMixin,
    VariablePartitionedTSMixin,
    VariablesEndpointMixin,
)
from meteora.utils import DateTimeType, KwargsType, VariablesType

# API endpoints
BASE_URL = "https://api.meteo.cat/xema/v1"
STATIONS_ENDPOINT = f"{BASE_URL}/estacions/metadades"
VARIABLES_ENDPOINT = f"{BASE_URL}/variables/mesurades/metadades"
TS_ENDPOINT = (
    BASE_URL
    + "/variables/mesurades/{variable_id}"
    + "/{period.year}/{period.month:02d}/{period.day:02d}"
)

# useful constants
STATIONS_GDF_ID_COL = "codi"
TS_DF_STATIONS_ID_COL = "codi"
TS_DF_TIME_COL = "data"
VARIABLES_ID_COL = "codi"
ECV_DICT = {
    # precipitation
    settings.ECV_PRECIPITATION: 35,  # "Precipitació",
    # pressure
    settings.ECV_PRESSURE: 34,  # "Pressió atmosfèrica",
    # radiation budget
    settings.ECV_RADIATION_SHORTWAVE: 39,  # "Radiació UV",
    # temperature
    settings.ECV_TEMPERATURE: 32,  # "Temperatura",
    # water vapour
    settings.ECV_RELATIVE_HUMIDITY: 33,  # "Humitat relativa",
    # wind
    settings.ECV_WIND_SPEED: 46,  # "Velocitat del vent a 2 m (esc.)",
    settings.ECV_WIND_DIRECTION: 47,  # "Direcció de vent 2 m (m. 1)",
}


class MeteocatClient(
    VariablePartitionedTSMixin,
    TimePartitionedTSMixin,
    APIKeyHeaderMixin,
    StationsEndpointMixin,
    VariablesEndpointMixin,
    BaseJSONClient,
):
    """Meteocat client.

    Parameters
    ----------
    region : str, Sequence, GeoSeries, GeoDataFrame, PathLike, or IO
        The region to process. This can be either:

        -  A string with a place name (Nominatim query) to geocode.
        -  A sequence with the west, south, east and north bounds.
        -  A geometric object, e.g., shapely geometry, or a sequence of geometric
           objects. In such a case, the value will be passed as the `data` argument of
           the GeoSeries constructor, and needs to be in the same CRS as the one used by
           the client's class (i.e., the `CRS` class attribute).
        -  A geopandas geo-series or geo-data frame.
        -  A filename or URL, a file-like object opened in binary ('rb') mode, or a Path
           object that will be passed to `geopandas.read_file`.
    api_key : str
        Meteocat API key.
    progress : bool, optional
        Whether to show a tqdm progress bar for partitioned time series fetches. If
        None, the value from `settings.SHOW_PROGRESS` is used.
    sjoin_kwargs : dict, optional
        Keyword arguments to pass to the `geopandas.sjoin` function when filtering the
        stations within the region. If None, the value from `settings.SJOIN_KWARGS` is
        used.
    """

    # geom constants
    X_COL = "coordenades.longitud"
    Y_COL = "coordenades.latitud"
    CRS = utils.LONLAT_CRS

    # time partition frequency
    _time_partition_freq = "D"

    # API endpoints
    _stations_endpoint = STATIONS_ENDPOINT
    _variables_endpoint = VARIABLES_ENDPOINT
    _ts_endpoint = TS_ENDPOINT

    # data frame labels constants
    _stations_gdf_id_col = STATIONS_GDF_ID_COL
    _ts_df_stations_id_col = TS_DF_STATIONS_ID_COL
    _ts_df_time_col = TS_DF_TIME_COL
    _variables_id_col = VARIABLES_ID_COL
    _ecv_dict = ECV_DICT

    def __init__(
        self,
        region: RegionType,
        api_key: str,
        *,
        progress: bool | None = None,
        **sjoin_kwargs: KwargsType,
    ) -> None:
        """Initialize Meteocat client."""
        self.region = region
        self._api_key = api_key
        if not sjoin_kwargs:
            sjoin_kwargs = settings.SJOIN_KWARGS.copy()
        self.SJOIN_KWARGS = sjoin_kwargs

        # need to call super().__init__() to set the cache
        super().__init__()
        if progress is not None:
            self.progress = progress

    def _ts_query_params(self, ts_params: Mapping) -> Mapping:
        return {}

    def _stations_df_from_content(self, response_content: Mapping) -> pd.DataFrame:
        return pd.json_normalize(response_content)

    def _variables_df_from_content(self, response_content: Mapping) -> pd.DataFrame:
        return pd.json_normalize(response_content)

    def _ts_df_from_content(self, response_content: Mapping):
        # process response
        response_df = pd.json_normalize(response_content)
        # filter stations
        response_df = response_df[response_df["codi"].isin(self.stations_gdf.index)]
        # extract observed data
        ts_df = pd.concat(
            response_df.apply(
                lambda row: pd.DataFrame(row["variables"][0]["lectures"]), axis=1
            ).tolist()
        )
        # add the station id column matching the observations
        ts_df[self._ts_df_stations_id_col] = (
            response_df[self._ts_df_stations_id_col]
            .repeat(
                response_df.apply(
                    lambda row: len(row["variables"][0]["lectures"]), axis=1
                )
            )
            .values
        )
        values_col = "valor"
        return ts_df.assign(
            **{self._ts_df_time_col: pd.to_datetime(ts_df[self._ts_df_time_col])}
        ).set_index([self._ts_df_stations_id_col, self._ts_df_time_col])[values_col]

    def get_ts_df(
        self,
        variables: VariablesType,
        start: DateTimeType,
        end: DateTimeType,
    ) -> pd.DataFrame:
        """Get time series data frame.

        Parameters
        ----------
        variables : str, int or list-like of str or int
            Target variables, which can be either a Meteocat variable code (integer or
            string) or an essential climate variable (ECV) following the Meteora
            nomenclature (string).
        start, end : datetime-like, str, int, float
            Values representing the start and end of the requested data period
            respectively. Accepts any datetime-like object that can be passed to
            pandas.Timestamp.

        Returns
        -------
        ts_df : pandas.DataFrame
            Long form data frame with a time series of measurements (second-level index)
            at each station (first-level index) for each variable (column).
        """
        ts_df = self._get_ts_df(
            variables,
            start=start,
            end=end,
        )
        units_map = ts_df.attrs.get("units")
        # filter time range to avoid including a full day after
        time_ser = ts_df.index.get_level_values(settings.TIME_COL).to_series()
        tz = time_ser.dt.tz
        ts_df = ts_df.loc[
            (
                slice(None),
                time_ser.between(
                    pd.Timestamp(start, tz=tz),
                    pd.Timestamp(end, tz=tz),
                    inclusive="both",
                ),
            ),
            :,
        ]
        if isinstance(units_map, Mapping):
            ts_df.attrs = ts_df.attrs.copy()
            ts_df.attrs["units"] = dict(units_map)
        return ts_df
