# (C) Copyright 2024 Anemoi contributors.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.


import logging
from typing import Any

from earthkit.data.utils.dates import to_datetime

from anemoi.inference.context import Context
from anemoi.inference.types import DataRequest
from anemoi.inference.types import Date
from anemoi.inference.types import ProcessorConfig
from anemoi.inference.types import State

from . import input_registry
from .grib import GribInput

LOG = logging.getLogger(__name__)


def rounded_area(area: list[float] | None) -> list[float] | None:
    """Round the area to a global extent if the surface is greater than 0.98.

    Parameters
    ----------
    area : Optional[List[float]]
        The area to be rounded.

    Returns
    -------
    Optional[List[float]]
        The rounded area or the original area if no rounding is needed.
    """
    try:
        surface = (area[0] - area[2]) * (area[3] - area[1]) / 180 / 360
        if surface > 0.98:
            return [90, 0.0, -90, 360]
    except TypeError:
        pass
    return area


def grid_is_valid(grid: str | list[float] | None) -> bool:
    """Check if the grid is valid.

    Parameters
    ----------
    grid : Optional[Union[str, List[float]]]
        The grid to be checked.

    Returns
    -------
    bool
        True if the grid is valid, False otherwise.
    """
    if grid is None:
        return False

    if isinstance(grid, str):
        return True

    try:
        [float(x) for x in grid]
        return True
    except TypeError:
        return False


def area_is_valid(area: list[float] | None) -> bool:
    """Check if the area is valid.

    Parameters
    ----------
    area : Optional[List[float]]
        The area to be checked.

    Returns
    -------
    bool
        True if the area is valid, False otherwise.
    """
    if area is None:
        return False

    if len(area) != 4:
        return False

    try:
        [float(x) for x in area]
        return True
    except TypeError:
        return False


def postproc(grid: str | list[float] | None, area: list[float] | None) -> dict[str, str | list[float]]:
    """Post-process the grid and area.

    Parameters
    ----------
    grid : Optional[Union[str, List[float]]]
        The grid to be post-processed.
    area : Optional[List[float]]
        The area to be post-processed.

    Returns
    -------
    Dict[str, Union[str, List[float]]]
        The post-processed grid and area.
    """
    pproc = dict()
    if grid_is_valid(grid):
        pproc["grid"] = grid

    if isinstance(area, str):
        area = [float(x) for x in area.split("/")]

    if area_is_valid(area):
        pproc["area"] = rounded_area(area)

    return pproc


def retrieve(
    requests: list[dict[str, Any]],
    grid: str | list[float] | None,
    area: list[float] | None,
    patch: Any | None = None,
    log: bool = True,
    **kwargs: Any,
) -> Any:
    """Retrieve data from MARS.

    Parameters
    ----------
    requests : List[Dict[str, Any]]
        The list of requests to be retrieved.
    grid : Optional[Union[str, List[float]]]
        The grid for the retrieval.
    area : Optional[List[float]]
        The area for the retrieval.
    patch : Optional[Any], optional
        Optional patch for the request, by default None.
    log : bool, optional
        Whether to log the requests, by default True.
    **kwargs : Any
        Additional keyword arguments.

    Returns
    -------
    Any
        The retrieved data.
    """
    import earthkit.data as ekd

    def _(r: DataRequest) -> str:
        mars = r.copy()
        for k, v in r.items():
            if isinstance(v, (list, tuple)):
                mars[k] = "/".join(str(x) for x in v)
            else:
                mars[k] = str(v)

        return ",".join(f"{k}={v}" for k, v in mars.items())

    pproc = postproc(grid, area)

    result = ekd.from_source("empty")
    for r in requests:
        if r.get("class") in ("rd", "ea"):
            r["class"] = "od"

        # ECMWF operational data has stream oper for 00 and 12 UTC and scda for 06 and 18 UTC

        if r.get("type") == "fc" and r.get("stream") == "oper" and r["time"] in ("0600", "1800"):
            r["stream"] = "scda"

        r.update(pproc)
        r.update(kwargs)

        if patch:
            r = patch(r)

        LOG.debug("%s", _(r))

        result += ekd.from_source("mars", r, log="default" if log else None)

    return result


@input_registry.register("mars")
class MarsInput(GribInput):
    """Get input fields from MARS."""

    trace_name = "mars"

    def __init__(
        self,
        context: Context,
        *,
        patches: list[tuple[dict[str, Any], dict[str, Any]]] | None = None,
        log: bool = True,
        variables: list[str] | None,
        pre_processors: list[ProcessorConfig] | None = None,
        namer: Any | None = None,
        purpose: str | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialize the MarsInput.

        Parameters
        ----------
        context : Any
            The context in which the input is used.
        namer : Optional[Any]
            Optional namer for the input.
        patches : Optional[List[Tuple[Dict[str, Any], Dict[str, Any]]]]
            Optional list of patches for the input.
        log : bool
            Whether to log the requests to MARS, by default True.
        **kwargs : Any
            Additional keyword to pass to the request to MARS.
        """
        super().__init__(
            context,
            variables=variables,
            pre_processors=pre_processors,
            purpose=purpose,
            namer=namer,
        )

        self.kwargs = kwargs
        self.patches = patches or []
        self.log = log

    def create_input_state(self, *, date: Date | None) -> State:
        """Create the input state for the given date.

        Parameters
        ----------
        date : Optional[Date]
            The date for which to create the input state.

        Returns
        -------
        State
            The created input state.
        """
        if date is None:
            date = to_datetime(-1)
            LOG.warning("MarsInput: `date` parameter not provided, using yesterday's date: %s", date)

        return self._create_input_state(
            self.retrieve(
                self.variables,
                [date + h for h in self.checkpoint.lagged],
            ),
            variables=self.variables,
            date=date,
        )

    def retrieve(self, variables: list[str], dates: list[Date]) -> Any:
        """Retrieve data for the given variables and dates.

        Parameters
        ----------
        variables : List[str]
            The list of variables to retrieve.
        dates : List[Any]
            The list of dates for which to retrieve the data.

        Returns
        -------
        Any
            The retrieved data.
        """
        requests = self.checkpoint.mars_requests(
            variables=variables,
            dates=dates,
            use_grib_paramid=self.context.use_grib_paramid,
            patch_request=self.patch_data_request,
        )

        if not requests:
            raise ValueError(f"No requests for {variables} ({dates})")

        kwargs = self.kwargs.copy()
        kwargs.setdefault("expver", "0001")
        kwargs.setdefault("grid", self.checkpoint.grid)
        kwargs.setdefault("area", self.checkpoint.area)

        return retrieve(
            requests,
            patch=self.patch,
            log=self.log,
            **kwargs,
        )

    def load_forcings_state(self, *, dates: list[Date], current_state: State) -> State:
        """Load the forcings state for the given variables and dates.

        Parameters
        ----------
        dates : List[Date]
            The list of dates for which to load the forcings state.
        current_state : State
            The current state to be updated with the loaded forcings state.

        Returns
        -------
        Any
            The loaded forcings state.
        """
        return self._load_forcings_state(
            self.retrieve(self.variables, dates),
            dates=dates,
            current_state=current_state,
        )

    def patch(self, request: dict[str, Any]) -> dict[str, Any]:
        """Patch the given request with predefined patches.

        Parameters
        ----------
        request : Dict[str, Any]
            The request to be patched.

        Returns
        -------
        Dict[str, Any]
            The patched request.
        """
        for match, keys in self.patches:
            if all(request.get(k) == v for k, v in match.items()):
                request.update(keys)

        return request
