"""Define models for existing (h52nx) context and helper functions."""

from __future__ import annotations

import re
from enum import Enum
from typing import Union
from pydantic import BaseModel, field_validator

"""Context can be given with settings keys to define a search context like
- 'bliss_tomo:sample_u' -> will use the dataset defined in 'bliss_tomo' to access sample_u information
- 'id16a:bliss_tomo:sample_u' -> If the dataset comes from id16a beamline will use the dataset defined in 'bliss_tomo' to access sample_u information
"""


class Proxy(Enum):
    """Determines if to access a dataset we read it directly or if we need to pass by a proxy like 'tomo_config'"""

    TOMO_CONFIG = "tomo_config"

    @classmethod
    def from_raw(cls, value: str):
        value = value.lower()
        if value == "tomoconfig":
            return Proxy.TOMO_CONFIG
        return cls(value)


class Beamline(Enum):
    id03 = "id03"
    bm05 = "bm05"
    bm18 = "bm18"
    id11 = "id11"
    id15 = "id15a"
    id16a = "id16a"
    id16b = "id16b"
    id17 = "id17"
    id19 = "id19"

    @classmethod
    def from_raw(cls, value: str):
        value = value.lower()
        value = value.replace("-", "")
        value = value.replace(":", "")
        value = value.replace("esrf", "")
        return cls(value)


_CONTEXT = Union[Proxy, Beamline]
_CONTEXT_CLASSES = {Proxy, Beamline}


class _Key(BaseModel):
    dataset: str
    contexts: set[_CONTEXT]

    @field_validator(
        "dataset",
        mode="plain",
    )
    @classmethod
    def check_dataset(cls, value: str) -> str:
        if ":" in value:
            raise ValueError(
                f"dataset name ({value}) is not expected to contain ':' char. Used as separator"
            )
        return value

    @staticmethod
    def resolve_dataset_operations(dataset_name: str) -> tuple[str, ...]:
        """Splits the dataset name to a chain of ('raw_dataset', operation,) ..."""
        # Use regex to split on operators (+, -, *, /) while keeping the operators
        tokens = re.split(r"([+\-*/])", dataset_name)
        # Remove empty strings and strip whitespace
        tokens = [token.strip() for token in tokens if token.strip()]
        return tokens


def extract_contexts(key: str, separator=":") -> _Key:
    """
    Split a given string into a dataset name and a set of context (beamline number, 'bliss_tomo'...)
    Order of the context doesn't matter but the dataset name must be at the end.
    """
    parts = key.split(separator)
    dataset_name = parts[-1]
    if len(parts) > 1:

        contexts = {_str_to_context(part) for part in parts[:-1]}
    else:
        contexts = set()
    if not dataset_name:
        raise ValueError(f"No dataset name provided in {key}")
    return _Key(dataset=dataset_name, contexts=contexts)


def _str_to_context(ctx: str) -> _CONTEXT:
    """convert a given string like 'ESRF-id11', 'tomo_config' to a context. Else raise an error"""
    for context in _CONTEXT_CLASSES:
        try:
            return context.from_raw(ctx)
        except ValueError:
            pass

    else:
        raise ValueError(f"Unrecognized context: {ctx}")


def key_to_str(key: _Key) -> str:
    """serialize instance of '_Key' to a string"""
    if not isinstance(key, _Key):
        raise TypeError(f"key is expected to be an instance of {_Key}")
    parts = [
        ctx.value
        for ctx in sorted(key.contexts, key=lambda context: isinstance(context, Proxy))
    ] + [
        key.dataset,
    ]
    return ":".join(parts)
