"""Define existing (h52nx) contexts and functions to handle them."""

from __future__ import annotations

import logging
import h5py
import numpy
import pint
from typing import TYPE_CHECKING, Callable
from silx.io.utils import h5py_read_dataset

from nxtomomill.models.h52nx import H52nxModel
from nxtomomill.models.h52nx._context import _Key, Proxy, Beamline as _BeamlineContext
from nxtomomill.utils.hdf5 import get_dataset_unit

_logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from nxtomomill.converter.hdf5.acquisition.blisstomoconfig import TomoConfig


_MODEL_FIELDS_WITH_KEYS = tuple(
    {
        key: value
        for key, value in H52nxModel.model_fields.items()
        if value.annotation == tuple[_Key, ...]
    }.keys()
)
"""Keys that can be contain a context == field that allows Field as type"""


def update_config_keys_from_tomo_config_context(
    configuration: H52nxModel, tomo_config: TomoConfig
):
    """Replace the tomo_config:xxx config by the dataset name contained in 'tomo_config'"""

    def map_dataset(dataset_name: str) -> str:
        """map a dataset from tomoconfig to it's raw name.

        ..warning:: tomoconfig dataset can store several dataset (alias, raw motor name). We are interested on the last only.
        """
        datasets = getattr(tomo_config, dataset_name)
        if datasets is None:
            raise KeyError(f"Cannot found tomo_config key: `{dataset_name}`")
        assert isinstance(
            datasets, (tuple, list)
        ), f"dataset is expected to be a tuple. Got {type(datasets)}"
        return datasets[-1]

    def apply_tomo_config_mapping(key: _Key) -> tuple[_Key, ...]:
        """
        Apply proxy / mapping from based on values contained in tomo_config.
        Example: 'sample_y+translation_y' becomes 'sy+ty' if tomo_config.sample_y == sy and tomo_config.translation_y == ty
        """
        split_dataset_and_operations = _Key.resolve_dataset_operations(key.dataset)
        mapping = [
            map_dataset(split_dataset_and_operations[0]),
        ]
        if len(split_dataset_and_operations) > 1:
            operators = split_dataset_and_operations[1::2]
            datasets = split_dataset_and_operations[2::2]
            for operation, dataset in zip(operators, datasets):
                mapping.extend((operation, map_dataset(dataset_name=dataset)))
        return "".join(mapping)

    def update_key(key: _Key) -> tuple[_Key, ...]:
        """One key return the original keys or two keys from tomoconfig: alias and motor"""
        if Proxy.TOMO_CONFIG in key.contexts:
            try:
                real_dataset_name = apply_tomo_config_mapping(key=key)
            except KeyError as e:
                _logger.error(e)
                # if the mapping cannot done with tomo_config log an error and remove the key from the list: consider it treated and invalid.
                return ()
            else:
                return (
                    _Key(
                        dataset=real_dataset_name,
                        contexts=key.contexts - {Proxy.TOMO_CONFIG},
                    ),
                )
        return (key,)

    update_configuration_keys(configuration, update_key)


def remove_config_keys_with_tomo_config_context(configuration: H52nxModel):
    """Remove all keys with a 'tomoconfig' context"""

    def update_key(key: _Key) -> tuple[_Key | None, ...]:
        if Proxy.TOMO_CONFIG in key.contexts:
            return (None,)
        else:
            return (key,)

    update_configuration_keys(configuration, update_key)


def update_config_keys_from_beamline_context(
    configuration: H52nxModel, beamline: _BeamlineContext | None
):
    """Function to filter settings according to the beamline"""

    def update_key(key: _Key) -> tuple[_Key | None, ...]:
        if not has_beamline_context(key) or (
            has_beamline_context(key) and beamline in key.contexts
        ):
            return (key,)
        else:
            return (None,)

    update_configuration_keys(configuration, update_key)


def has_beamline_context(key: _Key) -> bool:
    """Check if some beamline context are defined within the key"""
    for context in key.contexts:
        if context in _BeamlineContext:
            return True
    return False


def update_configuration_keys(
    configuration: H52nxModel,
    key_transform: Callable[[_Key]],
    fields: tuple[str, ...] = _MODEL_FIELDS_WITH_KEYS,
) -> None:
    """
    Generalized function to update keys in configuration fields.
    """

    def update_keys(keys):
        new_keys = []
        for key in keys:
            result = key_transform(key)
            new_keys.extend(result)
        return tuple(filter(None, new_keys))

    for keys_as_str in fields:
        keys_values = getattr(configuration, keys_as_str)
        setattr(configuration, keys_as_str, update_keys(keys_values))


class MissingValue(ValueError):
    """Raised if a dataset contains 'DIS' - not saved"""


def read_key(
    key: str,
    node: h5py.Group,
    expected_unit: pint.Unit,
    info_retrieve: str,
) -> pint.Quantity:
    """
    Util function to read a dataset from it's name.
    The dataset can be a composition of datasets and operations like 'sx+tx-ax'.
    In this case we first read the three dataset then apply sx+tx and finally tx-ax

    .. note::

        Unit of each dataset will be read and take into account. Dataset can link to scalar values or 1D arrays.
        If arrays have incoherent size (some points not recorded for example) then longest datasets will be cropped to match shortest dataset. End of the dataset will be ignored.
        A log message will be added.

    .. note::

        key can contain one of the following operator:

        * '+' for addition between two datasets
        * '-' for subtraction between two datasets

    .. note::

        If dataset contains '*DIS*'. Can happen for old dataset if the dataset has not been recorded. Will raise a MissingValue.

    """
    parts = _Key.resolve_dataset_operations(key)
    dataset_names = parts[::2]
    if len(parts) > 1:
        operators = parts[1::2]

    dataset_values = {}
    dataset_units = {}
    for dataset_name in dataset_names:
        values = h5py_read_dataset(node[dataset_name])
        dataset_values[dataset_name] = values
        dataset_units[dataset_name] = get_dataset_unit(
            dataset=node[dataset_name],
            default=expected_unit,
            from_dataset=f"{node[dataset_name].name} (reading {info_retrieve})",
        )
        # skip values containing '*DIS*'
        if isinstance(values, str) and values == "*DIS*":
            raise MissingValue(f"dataset {dataset_name} has not been recorded")

    first_dataset = dataset_names[0]
    result_unit = dataset_units[first_dataset] or expected_unit
    result = dataset_values[first_dataset] * result_unit

    # Map operation strings to NumPy functions
    OPERATION_MAP = {
        "+": numpy.add,
        "-": numpy.subtract,
    }

    if len(dataset_names) > 1:
        for operator, dataset_name in zip(operators, dataset_names[1:]):
            dataset_unit = dataset_units[dataset_name] or expected_unit
            dataset_quantity = dataset_values[dataset_name] * dataset_unit
            if (
                not numpy.isscalar(result.magnitude)
                and not numpy.isscalar(dataset_quantity.magnitude)
                and len(result.magnitude) != len(dataset_quantity.magnitude)
            ):
                # crop dataset if have an incoherent size.
                _logger.warning(
                    f"Incoherent number of points founds for {info_retrieve} when accessing '{key}'."
                )
                nb_pts = min(len(result), len(dataset_quantity))
                result = result[:nb_pts]
                dataset_quantity = dataset_quantity[:nb_pts]
            result = OPERATION_MAP[operator](result, dataset_quantity)

    return result
