# (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 calendar
import datetime
import logging
import os
from pathlib import PurePath
from typing import TYPE_CHECKING
from typing import Any

import numpy as np
import zarr
from anemoi.utils.config import load_any_dict_format
from anemoi.utils.config import load_config as load_settings
from numpy.typing import NDArray

if TYPE_CHECKING:
    from .dataset import Dataset

LOG = logging.getLogger(__name__)


def load_config() -> dict[str, Any]:
    """Load the configuration settings.

    Returns
    -------
    Dict[str, Any]
        The configuration settings.
    """
    return load_settings(defaults={"datasets": {"named": {}, "path": []}})


def add_named_dataset(name: str, path: str, **kwargs: Any) -> None:
    """Add a named dataset to the configuration.

    Parameters
    ----------
    name : str
        The name of the dataset.
    path : str
        The path to the dataset.
    **kwargs : Any
        Additional arguments.
    """
    config = load_config()
    if name["datasets"]["named"]:
        raise ValueError(f"Dataset {name} already exists")

    config["datasets"]["named"][name] = path


def add_dataset_path(path: str) -> None:
    """Add a dataset path to the configuration.

    Parameters
    ----------
    path : str
        The path to add.
    """
    config = load_config()

    if path not in config["datasets"]["path"]:
        config["datasets"]["path"].append(path)


def round_datetime(d: np.datetime64, dates: NDArray[np.datetime64], up: bool) -> np.datetime64:
    """Round up (or down) a datetime to the nearest date in a list of dates.

    Parameters
    ----------
    d : np.datetime64
        The datetime to round.
    dates : NDArray[np.datetime64]
        The list of dates.
    up : bool
        Whether to round up.

    Returns
    -------
    np.datetime64
        The rounded datetime.
    """
    if dates is None or len(dates) == 0:
        return d

    for i, date in enumerate(dates):
        if date == d:
            return date
        if date > d:
            if up:
                return date
            if i > 0:
                return dates[i - 1]
            return date
    return dates[-1]


def _as_date(
    d: int | str | np.datetime64 | datetime.date,
    dates: NDArray[np.datetime64],
    last: bool,
    frequency: datetime.timedelta | None = None,
) -> np.datetime64:
    """Convert a date to a numpy datetime64 object, rounding to the nearest date in a list of dates.

    Parameters
    ----------
    d : Union[int, str, np.datetime64, datetime.date]
        The date to convert.
    dates : NDArray[np.datetime64]
        The list of dates.
    last : bool
        Whether to round to the last date.
    frequency : Optional[datetime.timedelta]
        The frequency of the dataset.

    Returns
    -------
    np.datetime64
        The converted date.
    """
    # WARNING,  datetime.datetime is a subclass of datetime.date
    # so we need to check for datetime.datetime first

    if isinstance(d, (np.datetime64, datetime.datetime)):
        d = round_datetime(d, dates, up=not last)
        return np.datetime64(d)

    if isinstance(d, datetime.date):
        d = d.year * 10_000 + d.month * 100 + d.day

    try:
        d = int(d)
    except (ValueError, TypeError):
        pass

    if isinstance(d, int):
        delta = frequency
        if delta is None:
            delta = np.timedelta64(1, "s")
        delta = np.timedelta64(delta, "s")

        if len(str(d)) == 4:
            year = d
            if last:
                year = year + 1
                npdate = np.datetime64(f"{year:04}-01-01T00:00:00")
                return _as_date(npdate - delta, dates, last, frequency)
            else:
                return _as_date(np.datetime64(f"{year:04}-01-01T00:00:00"), dates, last, frequency)

        if len(str(d)) == 6:
            year = d // 100
            month = d % 100
            if last:
                month = month + 1
                if month > 12:
                    month = 1
                    year += 1
                npdate = np.datetime64(f"{year:04}-{month:02}-01T00:00:00")
                return _as_date(npdate - delta, dates, last, frequency)
            else:
                return _as_date(np.datetime64(f"{year:04}-{month:02}-01T00:00:00"), dates, last, frequency)

        if len(str(d)) == 8:
            year = d // 10000
            month = (d % 10000) // 100
            day = d % 100
            if last:
                day = day + 1
                if day > calendar.monthrange(year, month)[1]:
                    day = 1
                    month += 1
                    if month > 12:
                        month = 1
                        year += 1
                npdate = np.datetime64(f"{year:04}-{month:02}-{day:02}T00:00:00")
                return _as_date(npdate - delta, dates, last, frequency)
            else:
                return _as_date(np.datetime64(f"{year:04}-{month:02}-{day:02}T00:00:00"), dates, last, frequency)

    if isinstance(d, str):

        def isfloat(s: str) -> bool:
            try:
                float(s)
                return True
            except ValueError:
                return False

        if d.endswith("%") and isfloat(d[:-1]):
            x = float(d[:-1])
            if not 0 <= x <= 100:
                raise ValueError(f"Invalid date: {d}")
            i_float = x * len(dates) / 100

            epsilon = 2 ** (-30)
            if len(dates) > 1 / epsilon:
                LOG.warning("Too many dates to use percentage, one date may be lost in rounding")

            if last:
                index = int(i_float + epsilon) - 1
            else:
                index = int(i_float - epsilon)
            index = max(0, min(len(dates) - 1, index))
            return dates[index]

        if "-" in d and ":" in d:
            date, time = d.replace(" ", "T").split("T")
            year, month, day = (int(_) for _ in date.split("-"))
            hour, minute, second = (int(_) for _ in time.split(":"))
            return _as_date(
                np.datetime64(f"{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}"),
                dates,
                last,
                frequency,
            )

        if "-" in d:
            assert ":" not in d
            bits = d.split("-")
            if len(bits) == 1:
                return _as_date(int(bits[0]), dates, last, frequency)

            if len(bits) == 2:
                return _as_date(int(bits[0]) * 100 + int(bits[1]), dates, last, frequency)

            if len(bits) == 3:
                return _as_date(int(bits[0]) * 10000 + int(bits[1]) * 100 + int(bits[2]), dates, last, frequency)

        if ":" in d:
            assert len(d) == 5
            hour, minute = d.split(":")
            assert minute == "00"
            assert not last
            first = dates[0].astype(object)
            year = first.year
            month = first.month
            day = first.day

            return _as_date(np.datetime64(f"{year:04}-{month:02}-{day:02}T{hour}:00:00"), dates, last, frequency)

    raise NotImplementedError(f"Unsupported date: {d} ({type(d)})")


def as_first_date(
    d: int | str | np.datetime64 | datetime.date,
    dates: NDArray[np.datetime64],
    frequency: datetime.timedelta | None = None,
) -> np.datetime64:
    """Convert a date to the first date in a list of dates.

    Parameters
    ----------
    d : Union[int, str, np.datetime64, datetime.date]
        The date to convert.
    dates : NDArray[np.datetime64]
        The list of dates.
    frequency : Optional[datetime.timedelta]
        The frequency of the dataset.

    Returns
    -------
    np.datetime64
        The first date.
    """
    return _as_date(d, dates, last=False, frequency=frequency)


def as_last_date(
    d: int | str | np.datetime64 | datetime.date,
    dates: NDArray[np.datetime64],
    frequency: datetime.timedelta | None = None,
) -> np.datetime64:
    """Convert a date to the last date in a list of dates.

    Parameters
    ----------
    d : Union[int, str, np.datetime64, datetime.date]
        The date to convert.
    dates : NDArray[np.datetime64]
        The list of dates.
    frequency : Optional[datetime.timedelta]
        The frequency of the dataset.

    Returns
    -------
    np.datetime64
        The last date.
    """
    return _as_date(d, dates, last=True, frequency=frequency)


def _concat_or_join(datasets: list["Dataset"], kwargs: dict[str, Any]) -> tuple["Dataset", dict[str, Any]]:
    """Concatenate or join datasets based on their date ranges.

    Parameters
    ----------
    datasets : List[Dataset]
        The list of datasets.
    kwargs : Dict[str, Any]
        Additional arguments.

    Returns
    -------
    tuple[Dataset, Dict[str, Any]]
        The concatenated or joined dataset and remaining arguments.
    """
    if "adjust" in kwargs:
        raise ValueError("Cannot use 'adjust' without specifying 'concat' or 'join'")
    datasets, kwargs = _auto_adjust(datasets, kwargs)

    # Study the dates
    ranges = [(d.dates[0].astype(object), d.dates[-1].astype(object)) for d in datasets]

    if len(set(ranges)) == 1:
        from .join import Join

        return Join(datasets)._overlay(), kwargs

    from .concat import Concat

    Concat.check_dataset_compatibility(datasets)

    return Concat(datasets), kwargs


def _open(a: str | PurePath | dict[str, Any] | list[Any] | tuple[Any, ...]) -> "Dataset":
    """Open a dataset from various input types.

    Parameters
    ----------
    a : Union[str, PurePath, Dict[str, Any], List[Any], tuple[Any, ...]]
        The input to open.

    Returns
    -------
    Dataset
        The opened dataset.
    """
    from .dataset import Dataset
    from .stores import Zarr
    from .stores import zarr_lookup

    if isinstance(a, str) and len(a.split(".")) in [2, 3]:

        metadata_path = os.path.join(a, "metadata.json")
        if os.path.exists(metadata_path):
            metadata = load_any_dict_format(metadata_path)
            if "backend" not in metadata:
                raise ValueError(f"Metadata for {a} does not contain 'backend' key")

            from anemoi.datasets.data.records import open_records_dataset

            return open_records_dataset(a, backend=metadata["backend"])

    if isinstance(a, Dataset):
        return a.mutate()

    if isinstance(a, zarr.hierarchy.Group):
        return Zarr(a).mutate()

    if isinstance(a, str):
        return Zarr(zarr_lookup(a)).mutate()

    if isinstance(a, PurePath):
        return _open(str(a)).mutate()

    if isinstance(a, dict):
        return _open_dataset(**a).mutate()

    if isinstance(a, (list, tuple)):
        return _open_dataset(*a).mutate()

    raise NotImplementedError(f"Unsupported argument: {type(a)}")


def _auto_adjust(
    datasets: list["Dataset"],
    kwargs: dict[str, Any],
    exclude: list[str] | None = None,
) -> tuple[list["Dataset"], dict[str, Any]]:
    """Automatically adjust datasets based on specified criteria.

    Parameters
    ----------
    datasets : List[Dataset]
        The list of datasets.
    kwargs : Dict[str, Any]
        Additional arguments.
    exclude : Optional[List[str]]
        List of keys to exclude from adjustment.

    Returns
    -------
    tuple[List[Dataset], Dict[str, Any]]
        The adjusted datasets and remaining arguments.
    """
    if "adjust" not in kwargs:
        return datasets, kwargs

    adjust_list = kwargs.pop("adjust")
    if not isinstance(adjust_list, (tuple, list)):
        adjust_list = [adjust_list]

    ALIASES = {
        "all": ["select", "frequency", "start", "end"],
        "dates": ["start", "end", "frequency"],
        "variables": ["select"],
    }

    adjust_set = set()

    for a in adjust_list:
        adjust_set.update(ALIASES.get(a, [a]))

    if exclude is not None:
        adjust_set -= set(exclude)

    extra = set(adjust_set) - set(ALIASES["all"])
    if extra:
        raise ValueError(f"Invalid adjust keys: {extra}")

    subset_kwargs = [{} for _ in datasets]

    if "select" in adjust_set:
        assert "select" not in kwargs, "Cannot use 'select' in adjust and kwargs"

        variables = None

        for d in datasets:
            if variables is None:
                variables = set(d.variables)
            else:
                variables &= set(d.variables)

        if len(variables) == 0:
            raise ValueError("No common variables")

        for i, d in enumerate(datasets):
            if set(d.variables) != variables:
                subset_kwargs[i]["select"] = sorted(variables)

    if "start" or "end" in adjust_set:
        common = datasets[0].dates
        for d in datasets[0:]:
            common = np.intersect1d(common, d.dates)

    if "start" in adjust_set:
        assert "start" not in kwargs, "Cannot use 'start' in adjust and kwargs"
        start = min(common).astype(object)
        for i, d in enumerate(datasets):
            if start != d.dates[0]:
                subset_kwargs[i]["start"] = start

    if "end" in adjust_set:
        assert "end" not in kwargs, "Cannot use 'end' in adjust and kwargs"
        end = max(common).astype(object)
        for i, d in enumerate(datasets):
            if end != d.dates[-1]:
                subset_kwargs[i]["end"] = end

    if "frequency" in adjust_set:
        assert "frequency" not in kwargs, "Cannot use 'frequency' in adjust and kwargs"
        frequency = max(d.frequency for d in datasets)
        for i, d in enumerate(datasets):
            if d.frequency != frequency:
                subset_kwargs[i]["frequency"] = frequency

    datasets = [d._subset(**subset_kwargs[i]) for i, d in enumerate(datasets)]

    return datasets, kwargs


def _open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
    """Open a dataset.

    Parameters
    ----------
    *args : Any
        Positional arguments.
    **kwargs : Any
        Keyword arguments.

    Returns
    -------
    Dataset
        The opened dataset.
    """
    sets = []
    for a in args:
        sets.append(_open(a))

    if "observations" in kwargs:
        from .observations import observations_factory

        assert not sets, sets

        return observations_factory(args, kwargs).mutate()

    if "xy" in kwargs:
        # Experimental feature, may be removed
        from .xy import xy_factory

        assert not sets, sets
        return xy_factory(args, kwargs).mutate()

    if "x" in kwargs and "y" in kwargs:
        # Experimental feature, may be removed
        from .xy import xy_factory

        assert not sets, sets
        return xy_factory(args, kwargs).mutate()

    if "zip" in kwargs:
        # Experimental feature, may be removed
        from .xy import zip_factory

        assert not sets, sets
        return zip_factory(args, kwargs).mutate()

    if "chain" in kwargs:
        # Experimental feature, may be removed
        from .unchecked import chain_factory

        assert not sets, sets
        return chain_factory(args, kwargs).mutate()

    if "join" in kwargs:
        from .join import join_factory

        assert not sets, sets
        return join_factory(args, kwargs).mutate()

    if "concat" in kwargs:
        from .concat import concat_factory

        assert not sets, sets
        return concat_factory(args, kwargs).mutate()

    if "merge" in kwargs:
        from .merge import merge_factory

        assert not sets, sets
        return merge_factory(args, kwargs).mutate()

    if "ensemble" in kwargs:
        from .ensemble import ensemble_factory

        assert not sets, sets
        return ensemble_factory(args, kwargs).mutate()

    if "grids" in kwargs:
        from .grids import grids_factory

        assert not sets, sets
        return grids_factory(args, kwargs).mutate()

    if "cutout" in kwargs:
        from .grids import cutout_factory

        assert not sets, sets
        return cutout_factory(args, kwargs).mutate()

    if "complement" in kwargs:
        from .complement import complement_factory

        assert not sets, sets
        return complement_factory(args, kwargs).mutate()

    for name in ("datasets", "dataset"):
        if name in kwargs:
            datasets = kwargs.pop(name)
            if not isinstance(datasets, (list, tuple)):
                datasets = [datasets]
            for a in datasets:
                sets.append(_open(a))

    assert len(sets) > 0, (args, kwargs)

    if len(sets) > 1:
        dataset, kwargs = _concat_or_join(sets, kwargs)
        return dataset._subset(**kwargs)

    return sets[0]._subset(**kwargs)


def append_to_zarr(new_data: np.ndarray, new_dates: np.ndarray, zarr_path: str) -> None:
    """Append data from a subset (for one date) to the Zarr store.

    Parameters
    ----------
    new_data : np.ndarray
        The new data to append.
    new_dates : np.ndarray
        The new dates to append.
    zarr_path : str
        The path to the Zarr store.

    Notes
    -----
    - "dates" dataset is created with chunks equal to len(big_dataset.dates).
    - "data" dataset is created with chunk size 1 along the first (time) dimension.
    """
    print("Appending data for", new_dates, flush=True)
    # Re-open the zarr store to avoid root object accumulating memory.
    root = zarr.open(zarr_path, mode="a")
    # Convert new dates to strings (using str) regardless of input dtype.
    new_dates = np.array(new_dates, dtype="datetime64[s]")
    dates_ds = root["dates"]
    old_len = dates_ds.shape[0]
    dates_ds.resize((old_len + len(new_dates),))
    dates_ds[old_len:] = new_dates
    # Append to "data" dataset.
    data_ds = root["data"]
    old_shape = data_ds.shape  # (time, n_vars, ensembles, gridpoints)
    new_shape = (old_shape[0] + len(new_dates),) + old_shape[1:]
    data_ds.resize(new_shape)
    data_ds[old_shape[0] :] = new_data


def process_date(date: Any, big_dataset: "Dataset") -> tuple[np.ndarray, np.ndarray]:
    """Open the subset corresponding to the given date and return (date, subset).

    Parameters
    ----------
    date : Any
        The date to process.
    big_dataset : Dataset
        The dataset to process.

    Returns
    -------
    tuple[np.ndarray, np.ndarray]
        The subset and the date.
    """
    print("Processing:", date, flush=True)
    subset = _open_dataset(big_dataset, start=date, end=date).mutate()
    s = subset[:]
    date = subset.dates
    return s, date


def initialize_zarr_store(root: Any, big_dataset: "Dataset") -> None:
    """Initialize the Zarr store with the given dataset and recipe.

    Parameters
    ----------
    root : Any
        The root Zarr store.
    big_dataset : Dataset
        The dataset to initialize the store with.
    """
    ensembles = big_dataset.shape[2]
    # Create or append to "dates" dataset.
    if "dates" not in root:
        full_length = len(big_dataset.dates)
        root.create_dataset("dates", data=np.array([], dtype="datetime64[s]"), chunks=(full_length,))

    if "data" not in root:
        dims = (1, len(big_dataset.variables), ensembles, big_dataset.shape[-1])
        root.create_dataset(
            "data",
            shape=dims,
            dtype=np.float64,
            chunks=dims,
        )

    for k, v in big_dataset.statistics.items():
        if k not in root:
            root.create_dataset(
                k,
                data=v,
                compressor=None,
            )

    # Create spatial coordinate datasets if missing.
    if "latitudes" not in root or "longitudes" not in root:
        root.create_dataset("latitudes", data=big_dataset.latitudes, compressor=None)
        root.create_dataset("longitudes", data=big_dataset.longitudes, compressor=None)
    for k, v in big_dataset.metadata().items():
        if k not in root.attrs:
            root.attrs[k] = v
    # Set store-wide attributes if not already set.
    if "first_date" not in root.attrs:
        root.attrs["first_date"] = big_dataset.metadata()["start_date"]
        root.attrs["last_date"] = big_dataset.metadata()["end_date"]
        root.attrs["resolution"] = big_dataset.resolution
        root.attrs["name_to_index"] = {k: i for i, k in enumerate(big_dataset.variables)}
        root.attrs["ensemble_dimension"] = 2
        root.attrs["field_shape"] = big_dataset.field_shape
        root.attrs["flatten_grid"] = True
        root.attrs["recipe"] = {}


def _save_dataset(dataset: "Dataset", zarr_path: str, n_workers: int = 1) -> None:
    """Incrementally create (or update) a Zarr store from an Anemoi dataset.

    Parameters
    ----------
    dataset : Dataset
        anemoi-dataset opened from python to save to Zarr store
    zarr_path : str
        The path to the Zarr store.
    n_workers : int, optional
        The number of worker processes to use, by default 1.

    Notes
    -----
    Worker processes extract data for each date in parallel, but all writes
    to the store happen sequentially in the main process (i.e. single-writer).

    The "dates" dataset is created with chunking equal to the full length of
    big_dataset.dates, while "data" is chunked with 1 in the time dimension.
    """
    from concurrent.futures import ProcessPoolExecutor

    full_ds = dataset
    print("Opened full dataset.", flush=True)

    # Use ProcessPoolExecutor for parallel data extraction.
    # Workers return (date, subset) tuples.
    root = zarr.open(zarr_path, mode="a")
    initialize_zarr_store(root, full_ds)
    print("Zarr store initialised.", flush=True)

    existing_dates = np.array(sorted(root["dates"]), dtype="datetime64[s]")
    all_dates = full_ds.dates
    # To resume creation of the Zarr store in case the job is interrupted.
    dates_to_process = np.array(sorted(set(all_dates).difference(existing_dates)), dtype="datetime64[s]")

    use_pool = False

    if use_pool:
        with ProcessPoolExecutor(n_workers) as pool:
            futures = [pool.submit(process_date, date, full_ds) for date in dates_to_process]
            for future in futures:
                subset, date = future.result()
                # All appends happen sequentially here to
                #  avoid dates being added in a random order.
                append_to_zarr(subset, date, zarr_path)
    else:
        for date in dates_to_process:
            subset, date = process_date(date, full_ds)
            append_to_zarr(subset, date, zarr_path)
