# Copyright (c) 2018-2025 by xcube team and contributors
# Permissions are hereby granted under the terms of the MIT License:
# https://opensource.org/licenses/MIT.

import itertools
import os
import re
import uuid
import warnings
from collections.abc import Iterable, Mapping, Sequence
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import dask.array as da
import dask.array.core as dac
import distributed
import numpy as np

IntTuple = tuple[int, ...]
SliceTuple = tuple[slice, ...]
IntIterable = Iterable[int]
IntTupleIterable = Iterable[IntTuple]
SliceTupleIterable = Iterable[SliceTuple]

_CLUSTER_TAGS_ENV_VAR_NAME = "XCUBE_DASK_CLUSTER_TAGS"
_CLUSTER_ACCOUNT_ENV_VAR_NAME = "XCUBE_DASK_CLUSTER_ACCOUNT"


def compute_array_from_func(
    func: Callable[..., np.ndarray],
    shape: IntTuple,
    chunks: IntTuple,
    dtype: Any,
    name: str = None,
    ctx_arg_names: Sequence[str] = None,
    args: Sequence[Any] = None,
    kwargs: Mapping[str, Any] = None,
) -> da.Array:
    """Compute a dask array using the provided user function
    *func*, *shape*, and chunking *chunks*.

    The user function is expected to output the array's data
    blocks using arguments specified by *ctx_arg_names*, *args*,
    and *kwargs* and is expected to return a numpy array.

    You can request array and current block context information
    by specifying the optional *ctx_arg_names* keyword argument
    that is a sequence of names of special arguments passed to
    *user_func*. The following are available:

    * ``shape``: The array's shape. A tuple of ints.
    * ``chunks``: The array's chunks. A tuple of tuple of ints.
    * ``dtype``: The array's numpy data type.
    * ``name``: The array's name. A string or ``None``.
    * ``block_id``: The block's unique ID. An integer number
        ranging from zero to number of blocks minus one.
    * ``block_index``: The block's index as a tuple of ints.
    * ``block_shape``: The block's shape as a tuple of ints.
    * ``block_slices``: The block's shape as a tuple of int pair tuples.

    Args:
        func: User function that is called for each block of the
            array using arguments specified by *ctx_arg_names*,
            *args*, and *kwargs*. It must return a numpy array of
            shape "block_shape" and type *dtype*.
        shape: The array's shape. A tuple of sizes for each
            array dimension.
        chunks: The array's chunking. A tuple of chunk sizes for
            each array dimension. Must be of same length as *shape*.
        dtype: The array's numpy data type.
        name: The array's name.
        ctx_arg_names: Sequence names of arguments that are passed
            before *args* and *kwargs* to the user function.
        args: Arguments passed to the user function.
        kwargs: Keyword-arguments passed to the user function.

    Returns: A chunked dask array.
    """
    ctx_arg_names = ctx_arg_names or []
    args = args or []
    kwargs = kwargs or {}

    chunk_sizes = tuple(get_chunk_sizes(shape, chunks))
    chunk_counts = tuple(get_chunk_counts(shape, chunks))
    block_indexes, block_shapes, block_slices = get_block_iterators(chunk_sizes)

    ctx_values = dict(
        shape=tuple(shape),
        chunks=chunk_sizes,
        dtype=dtype,
        name=name,
    )

    blocks = _NestedList(shape=chunk_counts)
    block_id = 0
    for chunk_index, chunk_shape, block_slices in zip(
        block_indexes, block_shapes, block_slices
    ):
        ctx_values.update(
            block_id=block_id,
            block_index=tuple(chunk_index),
            block_shape=tuple(chunk_shape),
            block_slices=tuple(
                (chunk_slice.start, chunk_slice.stop) for chunk_slice in block_slices
            ),
        )
        ctx_args = [ctx_values[ctx_arg_name] for ctx_arg_name in ctx_arg_names]
        block_id += 1

        # We use our own name here, because dac.from_func() tokenizes args which for some reason takes forever
        block = dac.from_func(
            func,
            shape=chunk_shape,
            dtype=dtype,
            name=f"rectify_{name}-{uuid.uuid4()}",
            args=(*ctx_args, *args),
            kwargs=kwargs,
        )

        blocks[chunk_index] = block

    return da.block(blocks.data)


def get_block_iterators(
    chunk_sizes: IntTupleIterable,
) -> tuple[IntTupleIterable, IntTupleIterable, SliceTupleIterable]:
    chunk_sizes = tuple(chunk_sizes)
    chunk_slices_tuples = get_chunk_slice_tuples(chunk_sizes)
    chunk_ranges = get_chunk_ranges(chunk_sizes)
    block_indexes = itertools.product(*chunk_ranges)
    block_shapes = itertools.product(*chunk_sizes)
    block_slices = itertools.product(*chunk_slices_tuples)
    return block_indexes, block_shapes, block_slices


def get_chunk_sizes(shape: IntTuple, chunks: IntTuple) -> IntTupleIterable:
    for s, c in zip(shape, chunks):
        n = s // c
        if n * c < s:
            yield (c,) * n + (s % c,)
        else:
            yield (c,) * n


def get_chunk_counts(shape: IntTuple, chunks: IntTuple) -> Iterable[int]:
    for s, c in zip(shape, chunks):
        yield (s + c - 1) // c


def get_chunk_ranges(chunk_size_tuples: IntTupleIterable) -> Iterable[range]:
    return (range(len(chunk_size_tuple)) for chunk_size_tuple in chunk_size_tuples)


def get_chunk_slice_tuples(chunk_size_tuples: IntTupleIterable) -> SliceTupleIterable:
    return (
        tuple(get_chunk_slices(chunk_size_tuple))
        for chunk_size_tuple in chunk_size_tuples
    )


def get_chunk_slices(chunk_sizes: Sequence[int]) -> Iterable[slice]:
    stop = 0
    for i in range(len(chunk_sizes)):
        start = stop
        stop = start + chunk_sizes[i]
        yield slice(start, stop)


def new_cluster(
    provider: str = "coiled",
    name: Optional[str] = None,
    software: Optional[str] = None,
    n_workers: int = 4,
    resource_tags: Optional[dict[str, str]] = None,
    account: str = None,
    region: str = "eu-central-1",
    **kwargs,
) -> distributed.deploy.Cluster:
    """Create a new Dask cluster.

    Cloud resource tags can be specified in an environment variable
    XCUBE_DASK_CLUSTER_TAGS in the format
    ``tag_1=value_1:tag_2=value_2:...:tag_n=value_n``. In case of
    conflicts, tags specified in ``resource_tags`` will override tags
    specified by the environment variable.

    The cluster provider account name can be specified in an environment
    variable ``XCUBE_DASK_CLUSTER_ACCOUNT``. If the ``account`` argument is
    given to ``new_cluster``, it will override the value from the environment
    variable.

    Args:
        provider: identifier of the provider to use. Currently, only
            'coiled' is supported.
        name: name to use as an identifier for the cluster
        software: identifier for the software environment to be used.
        n_workers: number of workers in the cluster
        resource_tags: tags to apply to the cloud resources forming the
            cluster
        account: cluster provider account name
        **kwargs: further named arguments will be passed on to the
            cluster creation function
        region: default region where workers of the cluster will be
            deployed set to eu-central-1
    """

    if resource_tags is None:
        resource_tags = {}
    if _CLUSTER_ACCOUNT_ENV_VAR_NAME in os.environ:
        account_from_env_var = os.environ[_CLUSTER_ACCOUNT_ENV_VAR_NAME]
    else:
        account_from_env_var = None
        warnings.warn(
            f"Environment variable {_CLUSTER_ACCOUNT_ENV_VAR_NAME}"
            f" not set; cluster account name may be incorrect."
        )

    cluster_account = (
        account
        if account is not None
        else account_from_env_var if account_from_env_var is not None else "bc"
    )

    if provider == "coiled":
        try:
            import coiled
        except ImportError as e:
            raise ImportError(
                f"provider 'coiled' requires package'coiled' to be installed"
            ) from e
        if software is None and "JUPYTER_IMAGE" in os.environ:
            # If the JUPYTER_IMAGE environment variable is set, we're
            # presumably in a Z2JH deployment and can base a
            # Coiled environment on the same image.
            # First we construct an identifier from the user image specifier.
            current_image = os.environ["JUPYTER_IMAGE"]
            software = re.sub(
                "[:.]",
                "-",
                re.search(r"/([^/]+)$", current_image).group(1),
            )
            # If the referenced software environment doesn't exist yet as a
            # Coiled environment, create it from the currently used image.
            available_environments = coiled.list_software_environments(
                account=account
            ).keys()
            if software not in available_environments:
                coiled.create_software_environment(
                    name=software, container=current_image
                )

        # If software is (still) None, Coiled will try to mirror the current
        # environment automagically.
        coiled_params = dict(
            n_workers=n_workers,
            environ=None,
            tags=_collate_cluster_resource_tags(resource_tags),
            account=cluster_account,
            name=name,
            software=software,
            use_best_zone=True,
            compute_purchase_option="spot_with_fallback",
            shutdown_on_close=True,
            region=region,
        )
        coiled_params.update(kwargs)

        return coiled.Cluster(**coiled_params)

    raise NotImplementedError(f"Unknown provider {provider!r}")


def _collate_cluster_resource_tags(extra_tags: dict[str, str]) -> dict[str, str]:
    fallback_tags = {
        "cost-center": "unknown",
        "environment": "dev",
        "creator": "auto",
        "purpose": "xcube dask cluster",
        "user": (
            os.environ.get("JUPYTERHUB_USER")  # JupyterHub
            or os.environ.get("USER")  # Unixes
            or os.environ.get("USERNAME")  # Windows
            or os.getlogin()
            or ""
        ),
    }
    if _CLUSTER_TAGS_ENV_VAR_NAME in os.environ:
        kvps = os.environ[_CLUSTER_TAGS_ENV_VAR_NAME].split(":")
        env_var_tags = {
            (parts := kvp.split("=", maxsplit=1))[0]: parts[1] for kvp in kvps
        }
    else:
        warnings.warn(
            f"Environment variable {_CLUSTER_TAGS_ENV_VAR_NAME}"
            f" not set; cluster resource tags may be missing."
        )
        env_var_tags = {}
    return fallback_tags | env_var_tags | extra_tags


class _NestedList:
    """Utility class whose instances are used as input to dask.block()."""

    def __init__(self, shape: Sequence[int], fill_value: Any = None):
        self._shape = tuple(shape)
        self._data = self._new_data(shape, len(shape), fill_value, 0)

    @classmethod
    def _new_data(
        cls, shape: Sequence[int], ndim: int, fill_value: Any, dim: int
    ) -> Union[list[list], list[Any]]:
        return [
            (
                cls._new_data(shape, ndim, fill_value, dim + 1)
                if dim < ndim - 1
                else fill_value
            )
            for _ in range(shape[dim])
        ]

    @property
    def shape(self) -> tuple[int, ...]:
        return self._shape

    @property
    def data(self) -> Union[list[list], list[Any]]:
        return self._data

    def __len__(self) -> int:
        return len(self._data)

    def __setitem__(self, index: Union[int, slice, tuple], value: Any):
        data = self._data
        if isinstance(index, tuple):
            n = len(index)
            for i in range(n - 1):
                data = data[index[i]]
            data[index[n - 1]] = value
        else:
            data[index] = value

    def __getitem__(self, index: Union[int, slice, tuple]) -> Any:
        data = self._data
        if isinstance(index, tuple):
            n = len(index)
            for i in range(n - 1):
                data = data[index[i]]
            return data[index[n - 1]]
        else:
            return data[index]
