# 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 fnmatch
import itertools
import os
import os.path
import warnings
from collections.abc import Collection, Mapping
from functools import cached_property
from typing import Any, Callable, Optional, Union

import numpy as np
import pandas as pd
import pyproj
import xarray as xr

from xcube.constants import LOG
from xcube.core.mldataset import (
    BaseMultiLevelDataset,
    IdentityMultiLevelDataset,
    MultiLevelDataset,
    augment_ml_dataset,
    open_ml_dataset_from_python_code,
)
from xcube.core.store import (
    DATASET_TYPE,
    MULTI_LEVEL_DATASET_TYPE,
    DatasetDescriptor,
    DataStoreConfig,
    DataStorePool,
)
from xcube.core.tile import (
    DEFAULT_CMAP_NAME,
    DEFAULT_CMAP_NORM,
    DEFAULT_VALUE_RANGE,
    get_var_cmap_params,
    get_var_valid_range,
)
from xcube.core.varexpr import split_var_assignment
from xcube.server.api import ApiError, Context, ServerConfig
from xcube.server.config import is_absolute_path
from xcube.util.assertions import assert_given, assert_instance
from xcube.util.cache import parse_mem_size
from xcube.util.cmaps import (
    ColormapRegistry,
    create_colormap_from_config,
    load_custom_colormap,
)
from xcube.webapi.common.context import ResourcesContext
from xcube.webapi.places import PlacesContext

COMPUTE_DATASET = "compute_dataset"
COMPUTE_VARIABLES = "compute_variables"

# We use tilde, because it is not a reserved URI characters
STORE_DS_ID_SEPARATOR = "~"
DATA_STORE_IDS_WARNING = ["stac"]
FS_TYPE_TO_PROTOCOL = {
    "local": "file",
    "obs": "s3",
    "file": "file",
    "s3": "s3",
    "memory": "memory",
}
NON_MEMORY_FILE_SYSTEMS = ["local", "obs", "file", "s3"]

ALL_PLACES = "all"

MultiLevelDatasetOpener = Callable[["DatasetsContext", ServerConfig], MultiLevelDataset]

DatasetConfig = Mapping[str, Any]


class DatasetsContext(ResourcesContext):
    def __init__(self, server_ctx: Context, ml_dataset_openers=None):
        super().__init__(server_ctx)
        # noinspection PyTypeChecker
        self._places_ctx: PlacesContext = server_ctx.get_api_ctx("places")
        assert isinstance(self._places_ctx, PlacesContext)
        self._ml_dataset_openers = ml_dataset_openers
        # cache for all dataset configs
        # contains tuples of form (MultiLevelDataset, dataset_config)
        self._dataset_cache = dict()
        (
            self._data_store_pool,
            self._dataset_configs,
            self.entrypoint_dataset_id,
            self._dataset_groups_config,
        ) = self._process_dataset_configs(self.config, self.base_dir)
        self._cm_styles, self._colormap_registry = self._get_cm_styles()

    def on_dispose(self):
        with self.rlock:
            # Close all datasets
            for ml_dataset, _ in self._dataset_cache.values():
                # noinspection PyBroadException
                try:
                    ml_dataset.close()
                except Exception:
                    pass
            # Clear all caches
            if self._dataset_cache:
                self._dataset_cache.clear()
            if self._data_store_pool:
                self._data_store_pool.remove_all_store_configs()
            self._dataset_configs = None

    @property
    def places_ctx(self):
        return self._places_ctx

    @property
    def dataset_cache(self) -> dict[str, tuple[MultiLevelDataset, DatasetConfig]]:
        return self._dataset_cache

    @cached_property
    def access_control(self) -> dict[str, Any]:
        return self.config.get("AccessControl", {})

    @cached_property
    def required_scopes(self) -> list[str]:
        return self.access_control.get("RequiredScopes", [])

    @property
    def colormap_registry(self) -> ColormapRegistry:
        return self._colormap_registry

    def get_required_dataset_scopes(self, dataset_config: DatasetConfig) -> set[str]:
        return self._get_required_scopes(
            dataset_config, "read:dataset", "Dataset", dataset_config["Identifier"]
        )

    def get_required_variable_scopes(
        self, dataset_config: DatasetConfig, var_name: str
    ) -> set[str]:
        return self._get_required_scopes(
            dataset_config, "read:variable", "Variable", var_name
        )

    def _get_required_scopes(
        self,
        dataset_config: DatasetConfig,
        base_scope: str,
        value_name: str,
        value: str,
    ) -> set[str]:
        required_global_scopes = set(self.required_scopes)
        required_dataset_scopes = set(
            dataset_config.get("AccessControl", {}).get("RequiredScopes", [])
        )
        if not required_global_scopes and not required_dataset_scopes:
            return required_global_scopes
        required_dataset_scopes = required_global_scopes.union(required_dataset_scopes)
        base_scope_prefix = base_scope + ":"
        required_dataset_scopes = {
            scope
            for scope in required_dataset_scopes
            if scope == base_scope or scope.startswith(base_scope_prefix)
        }
        pattern_scope = base_scope_prefix + "{" + value_name + "}"
        if pattern_scope in required_dataset_scopes:
            # Replace "{base_scope}:{value_name}" by "{base_scope}:{value}"
            required_dataset_scopes.remove(pattern_scope)
            required_dataset_scopes.add(base_scope_prefix + value)
        return required_dataset_scopes

    def get_ml_dataset(self, ds_id: str) -> MultiLevelDataset:
        ml_dataset, _ = self._get_dataset_entry(ds_id)
        return ml_dataset

    def set_ml_dataset(self, ml_dataset: MultiLevelDataset):
        self._set_dataset_entry(
            (ml_dataset, _new_dataset_config(Identifier=ml_dataset.ds_id, Hidden=True))
        )

    def add_dataset(
        self,
        dataset: Union[xr.Dataset, MultiLevelDataset],
        ds_id: Optional[str] = None,
        title: Optional[str] = None,
        description: Optional[str] = None,
        style: Optional[str] = None,
        color_mappings: dict[str, dict[str, Any]] = None,
    ) -> str:
        """Add a dataset to the configuration"""
        assert_instance(dataset, (xr.Dataset, MultiLevelDataset), "dataset")
        if isinstance(dataset, xr.Dataset):
            ml_dataset = BaseMultiLevelDataset(dataset, ds_id=ds_id)
        else:
            ml_dataset = dataset
            if not ds_id:
                _ds_id = ml_dataset.ds_id
                # TODO: Fixes https://github.com/xcube-dev/xcube/issues/1007, but
                #   this does not solve the likely root cause
                #   in xcube.core.mldataset.fs.FsMultiLevelDataset, see ctor.
                if "/" in _ds_id:
                    ds_id = _ds_id.replace("/", "-")
            if ds_id:
                ml_dataset = IdentityMultiLevelDataset(dataset, ds_id=ds_id)
            dataset = ml_dataset.base_dataset
        ds_id = ml_dataset.ds_id
        dataset_configs = self.get_dataset_configs()
        for index, dataset_config in enumerate(dataset_configs):
            if dataset_config["Identifier"] == ds_id:
                del dataset_configs[index]
                break
        dataset_config = _new_dataset_config(
            Identifier=ds_id, Title=title, Description=description
        )
        if style is not None:
            dataset_config.update(dict(Style=style))
        if color_mappings is not None:
            style = style or ds_id
            dataset_config.update(dict(Style=style))
            self._cm_styles[style] = color_mappings
        self._dataset_cache[ds_id] = ml_dataset, dataset_config
        self._dataset_configs.append(dataset_config)
        return ds_id

    def remove_dataset(self, ds_id: str):
        assert_instance(ds_id, str, "ds_id")
        assert_given(ds_id, "ds_id")
        if ds_id in self._dataset_cache:
            del self._dataset_cache[ds_id]
        self._dataset_configs = [
            dc for dc in self._dataset_configs if dc["Identifier"] != ds_id
        ]

    def add_ml_dataset(
        self,
        ml_dataset: MultiLevelDataset,
        ds_id: Optional[str] = None,
        title: Optional[str] = None,
        description: Optional[str] = None,
    ):
        if ds_id:
            ml_dataset.ds_id = ds_id
        else:
            ds_id = ml_dataset.ds_id
        self._set_dataset_entry(
            (
                ml_dataset,
                (
                    dict(
                        Identifier=ds_id,
                        Title=title,
                        Description=description,
                    )
                ),
            )
        )

    def get_dataset(
        self, ds_id: str, expected_var_names: Collection[str] = None
    ) -> xr.Dataset:
        ml_dataset, _ = self._get_dataset_entry(ds_id)
        dataset = ml_dataset.base_dataset
        if expected_var_names:
            for var_name in expected_var_names:
                if var_name not in dataset:
                    raise ApiError.NotFound(
                        f'Variable "{var_name}" not found in dataset "{ds_id}"'
                    )
        return dataset

    def get_time_series_dataset(self, ds_id: str, var_name: str = None) -> xr.Dataset:
        dataset_config = self.get_dataset_config(ds_id)
        ts_ds_name = dataset_config.get("TimeSeriesDataset", ds_id)
        try:
            # Try to get more efficient, time-chunked dataset
            return self.get_dataset(
                ts_ds_name, expected_var_names=[var_name] if var_name else None
            )
        except ApiError.NotFound:
            # This happens, if the dataset pointed to by 'TimeSeriesDataset'
            # does not contain the variable given by var_name.
            return self.get_dataset(
                ds_id, expected_var_names=[var_name] if var_name else None
            )

    def get_variable_for_z(
        self, ds_id: str, var_name: str, z_index: int
    ) -> xr.DataArray:
        ml_dataset = self.get_ml_dataset(ds_id)
        index = ml_dataset.num_levels - 1 - z_index
        if index < 0 or index >= ml_dataset.num_levels:
            raise ApiError.NotFound(
                f'Variable "{var_name}" has no z-index {z_index} in dataset "{ds_id}"'
            )
        dataset = ml_dataset.get_dataset(index)
        if var_name not in dataset:
            raise ApiError.NotFound(
                f'Variable "{var_name}" not found in dataset "{ds_id}"'
            )
        return dataset[var_name]

    @classmethod
    def _maybe_assign_store_instance_ids(
        cls,
        dataset_configs: list[dict[str, Any]],
        data_store_pool: DataStorePool,
        base_dir: str,
    ) -> None:
        assignable_dataset_configs = [
            dc
            for dc in dataset_configs
            if "StoreInstanceId" not in dc
            and dc.get("FileSystem", "file") in NON_MEMORY_FILE_SYSTEMS
        ]
        # split into sub-lists according to file system and
        # non-root store params
        config_lists = []
        for config in assignable_dataset_configs:
            store_params = cls._get_other_store_params_than_root(config)
            file_system = config.get("FileSystem", "file")
            appended = False
            for config_list in config_lists:
                if config_list[0] == file_system and config_list[1] == store_params:
                    config_list[2].append(config)
                    appended = True
                    break
            if not appended:
                config_lists.append((file_system, store_params, [config]))

        for file_system, store_params, config_list in config_lists:
            # Retrieve paths per configuration
            paths = [dc["Path"] for dc in config_list]
            # Determine common prefixes of paths (and call them roots)
            roots = _get_roots(paths)
            for root in roots:
                fs_protocol = FS_TYPE_TO_PROTOCOL.get(file_system, file_system)
                store_params_with_root = store_params.copy()
                if fs_protocol == "file" and not is_absolute_path(root):
                    store_params_with_root["root"] = f"{base_dir}/{root}"
                else:
                    store_params_with_root["root"] = root
                # See if there already is a store with this configuration
                data_store_config = DataStoreConfig(
                    store_id=fs_protocol, store_params=store_params_with_root
                )
                store_instance_id = data_store_pool.get_store_instance_id(
                    data_store_config
                )
                if not store_instance_id:
                    # Create new store with new unique store instance id
                    counter = 1
                    while data_store_pool.has_store_instance(
                        f"{fs_protocol}_{counter}"
                    ):
                        counter += 1
                    store_instance_id = f"{fs_protocol}_{counter}"
                    data_store_pool.add_store_config(
                        store_instance_id, data_store_config
                    )
                for config in config_list:
                    path = config["Path"]
                    if path.startswith(root) and config.get("StoreInstanceId") is None:
                        config["StoreInstanceId"] = store_instance_id
                        new_path = path[len(root) :]
                        while new_path.startswith("/") or new_path.startswith("\\"):
                            new_path = new_path[1:]
                        config["Path"] = new_path

    @classmethod
    def _get_other_store_params_than_root(cls, dataset_config: DatasetConfig) -> dict:
        file_system = FS_TYPE_TO_PROTOCOL.get(dataset_config.get("FileSystem", "file"))
        if file_system != "s3":
            return {}
        storage_options = dict()
        if "Anonymous" in dataset_config:
            storage_options["anon"] = dataset_config["Anonymous"]
        client_kwargs = dict()
        if "Endpoint" in dataset_config:
            client_kwargs["endpoint_url"] = dataset_config["Endpoint"]
        if "Region" in dataset_config:
            client_kwargs["region_name"] = dataset_config["Region"]
        storage_options["client_kwargs"] = client_kwargs
        store_params = dict(storage_options=storage_options)
        return store_params

    @classmethod
    def get_dataset_configs_from_stores(
        cls, data_store_pool: DataStorePool
    ) -> list[DatasetConfig]:
        all_dataset_configs: list[DatasetConfig] = []
        for store_instance_id in data_store_pool.store_instance_ids:
            LOG.info(f"Scanning store {store_instance_id!r}")
            data_store_config = data_store_pool.get_store_config(store_instance_id)

            # Note by forman: This iterator chaining is inefficient.
            # Preferably, we should offer
            #
            # store_dataset_ids = data_store.get_data_ids(
            #     data_type=(DATASET_TYPE, MULTI_LEVEL_DATASET_TYPE)
            # )
            #

            store_dataset_configs: list[ServerConfig] = data_store_config.user_data
            if store_dataset_configs:
                for store_dataset_config in store_dataset_configs:
                    dataset_id_pattern = store_dataset_config.get("Path", "*")
                    if _is_wildcard(dataset_id_pattern):
                        if store_instance_id in DATA_STORE_IDS_WARNING:
                            warnings.warn(
                                f"The data store with ID '{store_instance_id}' has "
                                "many data IDs. Using wildcard patterns to select "
                                "datasets may cause a long setup time of the server."
                            )
                        data_store = data_store_pool.get_store(store_instance_id)
                        store_dataset_ids = itertools.chain(
                            data_store.get_data_ids(data_type=DATASET_TYPE),
                            data_store.get_data_ids(data_type=MULTI_LEVEL_DATASET_TYPE),
                        )
                        for store_dataset_id in store_dataset_ids:
                            if fnmatch.fnmatch(store_dataset_id, dataset_id_pattern):
                                all_dataset_configs.append(
                                    _get_selected_dataset_config(
                                        store_dataset_id,
                                        store_instance_id,
                                        store_dataset_config,
                                    )
                                )
                    else:
                        all_dataset_configs.append(
                            _get_selected_dataset_config(
                                store_dataset_config["Path"],
                                store_instance_id,
                                store_dataset_config,
                            )
                        )

        # # Just for testing:
        # debug_file = 'all_dataset_configs.json'
        # with open(debug_file, 'w') as stream:
        #     json.dump(all_dataset_configs, stream)
        #     LOG.debug(f'Wrote file {debug_file!r}')

        return all_dataset_configs

    def new_dataset_metadata(
        self, store_instance_id: str, dataset_id: str
    ) -> Optional[DatasetDescriptor]:
        data_store = self._data_store_pool.get_store(store_instance_id)
        dataset_metadata = data_store.describe_data(dataset_id, data_type="dataset")
        if dataset_metadata.crs is not None:
            crs = pyproj.CRS.from_string(dataset_metadata.crs)
            if not crs.is_geographic:
                LOG.warning(
                    f"Ignoring dataset {dataset_id!r} from"
                    f" store instance {store_instance_id!r}"
                    f" because it uses a non-geographic CRS"
                )
                return None
        # noinspection PyTypeChecker
        return dataset_metadata

    def get_dataset_config(self, ds_id: str) -> DatasetConfig:
        dataset_configs = self.get_dataset_configs()
        dataset_config = next(
            (dsd for dsd in dataset_configs if dsd["Identifier"] == ds_id), None
        )
        if dataset_config is None:
            raise ApiError.NotFound(f'Dataset "{ds_id}" not found')
        return dataset_config

    def get_dataset_configs(self) -> list[DatasetConfig]:
        assert self._dataset_configs is not None
        return self._dataset_configs

    def get_dataset_groups_configs(self) -> list[DatasetConfig]:
        assert self._dataset_groups_config is not None
        return self._dataset_groups_config

    def get_entrypoint_dataset_id(self) -> str | None:
        if self.entrypoint_dataset_id:
            return self.entrypoint_dataset_id
        return None

    def get_data_store_pool(self) -> DataStorePool:
        assert self._data_store_pool is not None
        return self._data_store_pool

    @classmethod
    def _process_dataset_configs(
        cls, config: ServerConfig, base_dir: str
    ) -> tuple[DataStorePool, list[dict[str, Any]], str, list[dict[str, Any]]]:
        data_store_configs = config.get("DataStores", [])
        dataset_configs = config.get("Datasets", [])
        entrypoint_dataset_id = config.get("EntrypointDatasetId", "")
        dataset_groups_configs = config.get("DatasetGroups", [])

        data_store_pool = DataStorePool()
        for data_store_config_dict in data_store_configs:
            store_instance_id = data_store_config_dict.get("Identifier")
            store_id = data_store_config_dict.get("StoreId")
            store_params = data_store_config_dict.get("StoreParams", {})
            store_dataset_configs = data_store_config_dict.get("Datasets")
            store_config = DataStoreConfig(
                store_id, store_params=store_params, user_data=store_dataset_configs
            )
            data_store_pool.add_store_config(store_instance_id, store_config)
        dataset_configs = dataset_configs + cls.get_dataset_configs_from_stores(
            data_store_pool
        )
        # Allow dataset_configs to be writable, because
        # _maybe_assign_store_instance_ids() will change
        # entries:
        dataset_configs = [dict(c) for c in dataset_configs]
        cls._maybe_assign_store_instance_ids(dataset_configs, data_store_pool, base_dir)
        return (
            data_store_pool,
            dataset_configs,
            entrypoint_dataset_id,
            dataset_groups_configs,
        )

    def get_rgb_color_mapping(
        self, ds_id: str, norm_range: tuple[float, float] = (0.0, 1.0)
    ) -> tuple[list[Optional[str]], list[tuple[float, float]]]:
        var_names = [None, None, None]
        norm_ranges = [norm_range, norm_range, norm_range]
        color_mappings = self.get_color_mappings(ds_id)
        if color_mappings:
            rgb_mapping = color_mappings.get("rgb")
            if rgb_mapping:
                components = "Red", "Green", "Blue"
                for i in range(3):
                    component = components[i]
                    component_config = rgb_mapping.get(component, {})
                    var_name = component_config.get("Variable")
                    norm_vmin, norm_vmax = component_config.get(
                        "ValueRange", norm_range
                    )
                    var_names[i] = var_name
                    norm_ranges[i] = norm_vmin, norm_vmax
        return var_names, norm_ranges

    def get_color_mapping(
        self, ds_id: str, var_name: str
    ) -> tuple[str, str, tuple[float, float]]:
        if split_var_assignment(var_name)[1]:
            # var_name is an assignment expression, the result is unknown yet,
            # so we must return default values here.
            return DEFAULT_CMAP_NAME, DEFAULT_CMAP_NORM, DEFAULT_VALUE_RANGE

        cmap_name = None
        cmap_norm = None
        cmap_vmin, cmap_vmax = None, None
        color_mappings = self.get_color_mappings(ds_id)
        if color_mappings:
            color_mapping = color_mappings.get(var_name)
            if color_mapping:
                assert "ColorFile" not in color_mapping
                cmap_name = color_mapping.get("ColorBar")
                cmap_norm = color_mapping.get("Norm")
                cmap_vmin, cmap_vmax = color_mapping.get("ValueRange", (None, None))

        cmap_range = cmap_vmin, cmap_vmax
        if cmap_name and cmap_norm and None not in cmap_range:
            # noinspection PyTypeChecker
            return cmap_name, cmap_norm, cmap_range

        ds = self.get_dataset(ds_id, expected_var_names=[var_name])
        var = ds[var_name]
        valid_range = get_var_valid_range(var)
        return get_var_cmap_params(var, cmap_name, cmap_norm, cmap_range, valid_range)

    def _get_cm_styles(self) -> tuple[dict[str, Any], ColormapRegistry]:
        custom_colormaps = {}
        cm_styles = {}

        # go through CustomColorMaps
        for ctx_cmap in self.config.get("CustomColorMaps", []):
            custom_colormap, config_parse = create_colormap_from_config(ctx_cmap)
            custom_colormaps[custom_colormap.cm_name] = custom_colormap
        # go through Styles
        for style in self.config.get("Styles", []):
            style_id = style["Identifier"]
            color_mappings = dict()
            for var_name, color_mapping in style["ColorMappings"].items():
                if "ColorFile" in color_mapping:
                    custom_cmap_path = self.get_config_path(
                        color_mapping, "ColorMappings", path_entry_name="ColorFile"
                    )
                    custom_colormap = custom_colormaps.get(custom_cmap_path)
                    if custom_colormap is None:
                        custom_colormap = load_custom_colormap(custom_cmap_path)
                        if custom_colormap is not None:
                            custom_colormaps[custom_cmap_path] = custom_colormap
                    if custom_colormap is not None:
                        color_mappings[var_name] = {
                            "ColorBar": custom_colormap.cm_name,
                            "ValueRange": (
                                custom_colormap.norm.vmin,
                                custom_colormap.norm.vmax,
                            ),
                        }
                else:
                    if color_mapping.get("ColorBar") in custom_colormaps:
                        cmap = custom_colormaps[color_mapping["ColorBar"]]
                        if (
                            cmap.cm_type == "categorical"
                            and "ValueRange" in color_mapping
                        ):
                            LOG.warning(
                                f"Custom color map {color_mapping['ColorBar']!r} is "
                                "categorical. ValueRange is ignored."
                            )
                            color_mapping = dict(
                                ColorBar=color_mapping["ColorBar"],
                                ValueRange=[
                                    min(cmap.values),
                                    max(cmap.values),
                                ],
                            )
                        if "ValueRange" not in color_mapping:
                            color_mapping = dict(
                                ColorBar=color_mapping["ColorBar"],
                                ValueRange=[
                                    min(cmap.values),
                                    max(cmap.values),
                                ],
                            )
                    color_mappings[var_name] = dict(color_mapping)

            cm_styles[style_id] = color_mappings

        return cm_styles, ColormapRegistry(*custom_colormaps.values())

    def get_color_mappings(self, ds_id: str) -> Optional[dict[str, dict[str, Any]]]:
        dataset_config = self.get_dataset_config(ds_id)
        style_id = dataset_config.get("Style", "default")
        return self._cm_styles.get(style_id, {})

    def _get_dataset_entry(self, ds_id: str) -> tuple[MultiLevelDataset, ServerConfig]:
        if ds_id not in self._dataset_cache:
            with self.rlock:
                self._set_dataset_entry(self._create_dataset_entry(ds_id))
        return self._dataset_cache[ds_id]

    def _set_dataset_entry(
        self, dataset_entry: tuple[MultiLevelDataset, DatasetConfig]
    ):
        ml_dataset, dataset_config = dataset_entry
        self._dataset_cache[ml_dataset.ds_id] = ml_dataset, dataset_config

    def _create_dataset_entry(
        self, ds_id: str
    ) -> tuple[MultiLevelDataset, DatasetConfig]:
        dataset_config = self.get_dataset_config(ds_id)
        ml_dataset = self._open_ml_dataset(dataset_config)
        return ml_dataset, dataset_config

    def _open_ml_dataset(self, dataset_config: DatasetConfig) -> MultiLevelDataset:
        ds_id: str = dataset_config.get("Identifier")
        store_instance_id = dataset_config.get("StoreInstanceId")
        if store_instance_id:
            data_store_pool = self.get_data_store_pool()
            data_store = data_store_pool.get_store(store_instance_id)
            data_id = dataset_config.get("Path")
            open_params = dict(dataset_config.get("StoreOpenParams") or {})
            open_params_schema = data_store.get_open_data_params_schema(data_id=data_id)

            # Inject cache_size=chunk_cache_capacity, if given and possible
            chunk_cache_capacity = self.get_dataset_chunk_cache_capacity(dataset_config)
            if (
                chunk_cache_capacity
                and "cache_size" not in open_params
                and "cache_size" in open_params_schema.properties
            ):
                open_params["cache_size"] = chunk_cache_capacity

            # Inject data_type="mldataset", if possible
            if (
                "data_type" not in open_params
                and "data_type" in open_params_schema.properties
                and "mldataset" in data_store.get_data_types()
            ):
                open_params["data_type"] = "mldataset"

            with self.measure_time(
                tag=f"Opened dataset {ds_id!r} from data store {store_instance_id!r}"
            ):
                dataset = data_store.open_data(data_id, **open_params)

            if isinstance(dataset, MultiLevelDataset):
                # Expected, nominal case.
                ml_dataset: MultiLevelDataset = dataset
            else:
                # Fallback. Usually results in poor tile computation performance.
                ml_dataset = BaseMultiLevelDataset(dataset)

            ml_dataset.ds_id = ds_id

        else:
            fs_type = dataset_config.get("FileSystem")
            if fs_type != "memory":
                raise ApiError.InvalidServerConfig(
                    f"Invalid FileSystem {fs_type!r} in dataset configuration {ds_id!r}"
                )

            with self.measure_time(tag=f"Opened dataset {ds_id!r} from {fs_type!r}"):
                ml_dataset = _open_ml_dataset_from_python_code(self, dataset_config)

        augmentation = dataset_config.get("Augmentation")
        if augmentation:
            script_path = self.get_config_path(
                augmentation, f"'Augmentation' of dataset configuration {ds_id}"
            )
            input_parameters = augmentation.get("InputParameters", {})
            input_parameters = self.eval_config_value(input_parameters)
            callable_name = augmentation.get("Class")
            is_factory = callable_name is not None
            if not is_factory:
                callable_name = augmentation.get("Function", COMPUTE_VARIABLES)
            ml_dataset = augment_ml_dataset(
                ml_dataset,
                script_path,
                callable_name,
                self.get_ml_dataset,
                self.set_ml_dataset,
                input_parameters=input_parameters,
                is_factory=is_factory,
                exception_type=ApiError.InvalidServerConfig,
            )
        return ml_dataset

    def get_legend_label(self, ds_id: str, var_name: str):
        dataset = self.get_dataset(ds_id)
        if var_name in dataset:
            ds = self.get_dataset(ds_id)
            units = ds[var_name].units
            return units
        raise ApiError.NotFound(f'Variable "{var_name}" not found in dataset "{ds_id}"')

    def get_dataset_place_groups(
        self, ds_id: str, base_url: str, load_features=False
    ) -> list[dict]:
        dataset_config = self.get_dataset_config(ds_id)

        place_group_id_prefix = f"DS-{ds_id}-"

        def predicate(pg_id: str, _: Any):
            return pg_id.startswith(place_group_id_prefix)

        place_groups = self.places_ctx.get_cached_place_groups(predicate=predicate)
        if place_groups:
            # for a given dataset, all place groups cached,
            # so the returned list is complete
            return place_groups

        place_groups = self.places_ctx.load_place_groups(
            dataset_config.get("PlaceGroups", []),
            base_url,
            is_global=False,
            load_features=load_features,
            qualifiers=[ds_id],
        )

        for place_group in place_groups:
            place_group_id = place_group_id_prefix + place_group["id"]
            self.places_ctx.set_cached_place_group(place_group_id, place_group)

        return place_groups

    def get_dataset_place_group(
        self, ds_id: str, place_group_id: str, base_url: str, load_features=False
    ) -> dict:
        place_groups = self.get_dataset_place_groups(
            ds_id, base_url, load_features=False
        )
        for place_group in place_groups:
            if place_group_id == place_group["id"]:
                if load_features:
                    self.places_ctx.load_place_group_features(place_group)
                return place_group
        raise ApiError.NotFound(f'Place group "{place_group_id}" not found')

    def get_dataset_and_coord_variable(self, ds_name: str, dim_name: str):
        ds = self.get_dataset(ds_name)
        if dim_name not in ds.coords:
            raise ApiError.NotFound(
                f"Dimension {dim_name!r} has no coordinates in dataset {ds_name!r}"
            )
        return ds, ds.coords[dim_name]

    @classmethod
    def get_var_indexers(
        cls,
        ds_name: str,
        var_name: str,
        var: xr.DataArray,
        dim_names: list[str],
        query_params: dict[str, str],
    ) -> dict[str, Any]:
        var_indexers = dict()
        for dim_name in dim_names:
            if dim_name not in var.coords:
                raise ApiError.BadRequest(
                    f"Dimension {dim_name!r} of variable {var_name!r}"
                    f" in dataset {ds_name!r} has no coordinates"
                )
            coord_var = var.coords[dim_name]
            dim_value_str = query_params.get(dim_name)
            try:
                if dim_value_str is None:
                    var_indexers[dim_name] = coord_var.values[0]
                elif dim_value_str == "current":
                    var_indexers[dim_name] = coord_var.values[-1]
                elif np.issubdtype(coord_var.dtype, np.floating):
                    var_indexers[dim_name] = float(dim_value_str)
                elif np.issubdtype(coord_var.dtype, np.integer):
                    var_indexers[dim_name] = int(dim_value_str)
                elif np.issubdtype(coord_var.dtype, np.datetime64):
                    if "/" in dim_value_str:
                        date_str_1, date_str_2 = dim_value_str.split("/", maxsplit=1)
                        var_indexer_1 = pd.to_datetime(date_str_1)
                        var_indexer_2 = pd.to_datetime(date_str_2)
                        var_indexers[dim_name] = (
                            var_indexer_1 + (var_indexer_2 - var_indexer_1) / 2
                        )
                    else:
                        date_str = dim_value_str
                        var_indexers[dim_name] = pd.to_datetime(date_str)
                else:
                    raise ValueError(
                        f"unable to convert value"
                        f" {dim_value_str!r} to {coord_var.dtype!r}"
                    )
            except ValueError as e:
                raise ApiError.BadRequest(
                    f"{dim_value_str!r} is not a valid value"
                    f" for dimension {dim_name!r} "
                    f"of variable {var_name!r} of dataset {ds_name!r}"
                ) from e
        return var_indexers

    def get_dataset_chunk_cache_capacity(
        self, dataset_config: DatasetConfig
    ) -> Optional[int]:
        cache_size = self.get_chunk_cache_capacity(dataset_config, "ChunkCacheSize")
        if cache_size is None:
            cache_size = self.get_chunk_cache_capacity(
                self.config, "DatasetChunkCacheSize"
            )
        return cache_size

    @classmethod
    def get_chunk_cache_capacity(
        cls, config: Mapping[str, Any], cache_size_key: str
    ) -> Optional[int]:
        cache_size = config.get(cache_size_key, None)
        if not cache_size:
            return None
        elif isinstance(cache_size, str):
            try:
                cache_size = parse_mem_size(cache_size)
            except ValueError:
                raise ApiError.InvalidServerConfig(f"Invalid {cache_size_key}")
        elif not isinstance(cache_size, int) or cache_size < 0:
            raise ApiError.InvalidServerConfig(f"Invalid {cache_size_key}")
        return cache_size


def _open_ml_dataset_from_python_code(
    ctx: DatasetsContext, dataset_config: DatasetConfig
) -> MultiLevelDataset:
    ds_id = dataset_config.get("Identifier")
    path = ctx.get_config_path(dataset_config, f"dataset configuration {ds_id}")
    callable_name = dataset_config.get("Class")
    is_factory = callable_name is not None
    if not is_factory:
        callable_name = dataset_config.get("Function", COMPUTE_DATASET)
    input_dataset_ids = dataset_config.get("InputDatasets", [])
    input_parameters = dataset_config.get("InputParameters", {})
    input_parameters = ctx.eval_config_value(input_parameters)
    chunk_cache_capacity = ctx.get_dataset_chunk_cache_capacity(dataset_config)
    if chunk_cache_capacity:
        warnings.warn(
            "chunk cache size is not effective for datasets computed from scripts"
        )
    for input_dataset_id in input_dataset_ids:
        if not ctx.get_dataset_config(input_dataset_id):
            raise ApiError.InvalidServerConfig(
                f"Invalid dataset configuration {ds_id!r}:"
                f" Input dataset {input_dataset_id!r} of"
                f" callable {callable_name!r}"
                f" must reference another dataset"
            )
    return open_ml_dataset_from_python_code(
        path,
        callable_name=callable_name,
        input_ml_dataset_ids=input_dataset_ids,
        input_ml_dataset_getter=ctx.get_ml_dataset,
        input_parameters=input_parameters,
        is_factory=is_factory,
        ds_id=ds_id,
        exception_type=ApiError.InvalidServerConfig,
    )


def _is_not_empty(prefix):
    return prefix != "" and prefix != "/" and prefix != "\\"


def _get_roots(paths):
    list.sort(paths)
    prefixes = _get_common_prefixes(paths)
    if len(prefixes) < 1:
        roots = [""]
    else:
        # perform further step to merge prefixes with same start
        prefixes = list(set(prefixes))
        prefixes.sort()
        roots = []
        root_candidate = prefixes[0]
        for root in prefixes[1:]:
            common_root = os.path.commonprefix([root_candidate, root])
            if _is_not_empty(common_root):
                root_candidate = common_root
            else:
                roots.append(root_candidate)
                root_candidate = root
        roots.append(root_candidate)
    if roots[0] == "":
        roots.append(roots.pop(0))
    return roots


def _get_common_prefixes(p):
    # Recursively examine a list of paths for common prefixes:
    # If no common prefix is found, split the list in half and
    # examine each half separately
    prefix = os.path.commonprefix(p)
    # ensure prefix does not end with full or partial directory
    # or file name
    prefix = prefix[: max(_last_index(prefix, "/"), _last_index(prefix, "\\"), 0)]
    if _is_not_empty(prefix) or len(p) == 1:
        return [prefix]
    else:
        return _get_common_prefixes(p[: int(len(p) / 2)]) + _get_common_prefixes(
            p[int(len(p) / 2) :]
        )


def _get_selected_dataset_config(
    store_dataset_id: str, store_instance_id: str, dataset_config_base: dict
) -> dict:
    LOG.debug(f"Selected dataset {store_dataset_id!r}")
    dataset_config = dict(StoreInstanceId=store_instance_id, **dataset_config_base)
    if "Identifier" in dataset_config and dataset_config["Path"] != store_dataset_id:
        raise ApiError.InvalidServerConfig(
            "User-defined identifiers can only be assigned"
            " to datasets with non-wildcard paths."
        )
    elif "Identifier" not in dataset_config:
        dataset_config["Path"] = store_dataset_id
        dataset_config["Identifier"] = (
            f"{store_instance_id}{STORE_DS_ID_SEPARATOR}{store_dataset_id}"
        )
    return dataset_config


def _last_index(prefix, symbol):
    try:
        return prefix.rindex(symbol)
    except ValueError:
        return -1


_MULTI_LEVEL_DATASET_OPENERS = {
    "memory": _open_ml_dataset_from_python_code,
}


def _is_wildcard(string: str) -> bool:
    return "?" in string or "*" in string


def _new_dataset_config(**kwargs) -> dict[str, Any]:
    return {k: v for k, v in kwargs.items() if v is not None}
