# Copyright 2024-2025 Canonical Ltd.
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License version 3 as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Base classes and functions for composing Slurm data models."""

# pyright: reportIncompatibleMethodOverride=false
#   This override is required because `pyright` dislikes in `ModelMapping` how
#   `_MapModel` and `MutableSequence` have conflicting signatures for the
#   `update` method. Potentially a candidate for a later refactor.

__all__ = [
    "classproperty",
    "Metadata",
    "Mode",
    "Model",
    "ModelList",
    "ModelMapping",
    "make_model_builder",
]

import json
from abc import ABC, ABCMeta, abstractmethod
from collections import Counter, deque
from collections.abc import (
    Callable,
    Iterable,
    Iterator,
    Mapping,
    MutableMapping,
    MutableSequence,
    Sequence,
)
from dataclasses import dataclass
from enum import Enum, auto
from functools import singledispatchmethod
from inspect import isclass
from itertools import filterfalse, takewhile
from types import MethodType, NoneType, UnionType, get_original_bases
from typing import (
    Annotated,
    Any,
    TypeVar,
    cast,
    get_args,
    get_origin,
    get_type_hints,
)

from jsonschema import ValidationError, validate
from typing_extensions import Self

from ..exceptions import ModelError
from .callback import Callback, DefaultCallback


class classproperty(property):  # noqa N801
    """Class property attribute.

    Create a class-level, read-only property.
    """

    def __get__(self, __instance, __owner: Any | None = None) -> Any:  # noqa D105
        return super().__get__(__owner)


class subclassdispatch:  # noqa N801
    """Like `singledispatch`, but dispatch based on subclass instead.

    Dispatches functions based on what the passed argument is a subclass of.
    """

    def __init__(self, func: Callable[..., Any]) -> None:
        self._registry = {}
        self._default_func = func
        self.__name__ = func.__name__

    def register(self, func: Callable[..., Any]) -> None:
        _, cls = next(iter(get_type_hints(func).items()))
        self._registry[cls] = func

    def dispatch(self, sub_cls: Any) -> Callable[..., Any]:
        if not isinstance(sub_cls, type):
            sub_cls = type(sub_cls)

        for cls, func in self._registry.items():
            if issubclass(sub_cls, cls):
                return func

        return self._default_func

    def __call__(self, *args, **kwargs) -> Any:
        return self.dispatch(args[0])(*args, **kwargs)

    def __get__(self, instance, owner) -> Self | MethodType:
        if instance is not None:
            return MethodType(self, instance)
        else:
            return self


class Mode(Enum):
    r"""Mode to parse/marshal Slurm configuration data.

    Attributes:
        ONELINE: Parse/marshal Slurm configuration from/to one line.
        STANZA: Parse/marshal Slurm configuration from/to a stanza.

    Examples:
        If mode is set to `ONELINE`, parse/marshal using " " as the separator:
            parse:
                "key1=value key2=value" -> {"key1": "value", "key2": "value"}
            marshal:
                 {"key1": "value", "key2": "value"} -> "key1=value key2=value"

        If mode is set to `STANZA`, parse\\marshal using "\n" as the separator:
            parse:
                "key1=value\nkey2=value" -> {"key1": "value", "key2": "value"}
            marshal:
                {"key1": "value", "key2": "value"} -> "key1=value\nkey2=value"
    """

    ONELINE = auto()
    STANZA = auto()


@dataclass(frozen=True)
class Metadata:
    r"""Metadata for Slurm configuration field.

    Attributes:
        origin: Name of field in configuration model.
        alias: Alias(es) used to identify field in a configuration file.
        sep: Separator used to split a field's key and value.
        unique: Flag for if field can occur only once or multiple times in a configuration file.
        primary: Flag for if a field should be at the front of a line in marshalled output.
        callback: Callbacks for handling field during parsing and marshalling.
        default_factory:
            A zero argument callable that will provide a default value if either
            `del` is called on the field, or the field is not present in the configuration file.

    Methods:
        parse: Alias for accessing `callback.parser` method.
        marshal: Alias for accessing `callback.marshaller` method.
    """

    origin: str = ""
    alias: str = ""
    sep: str = "="
    unique: bool = True
    primary: bool = False
    callback: Callback = DefaultCallback
    default_factory: Callable[[], Any] | None = None

    def parse(self, expr: str) -> Any:
        """Alias for accessing `callback.parser` method."""
        return self.callback.parser(expr, self.sep)

    def marshal(self, key: str, value: str) -> str:
        """Alias for accessing `callback.marshaller` method."""
        return self.callback.marshaller(key, value, self.sep)


class _ModelMeta(ABCMeta):
    """Metaclass for defining configuration models."""

    def __new__(
        mcls,  # noqa N804
        name: str,
        bases: tuple[type, ...],
        namespace: dict[str, Any],
        /,
        **kwargs: dict[str, Any],
    ) -> "_ModelMeta":
        field: str

        if "__annotations__" in namespace:
            for field, annotation in namespace["__annotations__"].items():
                metadata = _get_metadata(annotation)
                namespace[field] = _make_property(_format_field(field), metadata.default_factory)

        return super().__new__(mcls, name, bases, namespace, **kwargs)


class _ModelBase(metaclass=_ModelMeta):
    """Helper class for providing base methods to models through inheritance."""

    def __init__(self, instance: Any, /) -> None:
        instance = self.validate_model(instance)
        if self.__model_builder__:
            object.__setattr__(self, "_model_data", self.__model_builder__(instance))
        else:
            object.__setattr__(self, "_model_data", instance)

    @classproperty
    @abstractmethod
    def __model_schema__(self) -> dict[str, Any]:  # noqa N805
        """Model json schema."""

    @classproperty
    def __model_builder__(cls) -> Callable[[Any], Any] | None:  # noqa N805
        """Model builder."""
        return None

    @classproperty
    def __model_mode__(cls) -> Mode:  # noqa N805
        """Model parse/marshal mode."""
        return Mode.STANZA

    @classmethod
    def from_str(cls, s: str, /) -> Self:
        """Create model from Slurm configuration string(s)."""
        return cls(_parse(cls, s))

    @classmethod
    def from_json(cls, obj: str | bytes | bytearray, /) -> Self:
        """Create model from a json object."""
        return cls(json.loads(obj))

    def json(self) -> str:
        """Get model as a json object."""
        return json.dumps(self.validate_model())

    def validate_model(self, instance: Iterable[Any] | None = None, /) -> Any:
        """Validate a model against its defined json schema.

        Args:
            instance:
                Optional dictionary object to validate against model json schema.
                If left set to `None`, this method will validate the attribute `_model_data`
                against the model json schema instead.

        Raises:
            ModelError:
                Raised if `instance` or the `_model_data` attribute fails to validate
                against the model's json schema.

        Returns:
            Expanded dictionary or list object. Sub-models must be rebuilt if validated output
            will be used for marshalling or overwriting the contents of the `_model_data`
            attribute. Output can be ignored if just performing a validation.
        """
        instance = (
            _get_as_builtin_object(instance)
            if instance is not None
            else _get_as_builtin_object(self._model_data)
        )
        try:
            validate(instance, schema=self.__model_schema__)
        except ValidationError as e:
            raise ModelError(e.message)

        return instance

    def __getattr__(self, item: str) -> Any:  # noqa D105
        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")

    def __setattr__(self, key: str, value: Any) -> Any:  # noqa D105
        if not hasattr(self, key):
            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{key}'")

        super().__setattr__(key, value)

    def __str__(self) -> str:  # noqa D105
        self.validate_model()
        return _marshal(self)


_TModelBase = TypeVar("_TModelBase", bound=_ModelBase)
_KT = TypeVar("_KT", bound=str)


class _MapModel(_ModelBase, ABC):
    """Helper class for providing mapping-related methods to models through inheritance."""

    def __init__(self, m: dict[str, Any] | None = None, /, **kwargs) -> None:
        super().__init__((m if m else {}) | kwargs)

    @classmethod
    def from_dict(cls, m: dict[str, Any], /) -> Self:
        """Create model from dictionary object."""
        return cls(m)

    def dict(self) -> dict[str, Any]:
        """Get model as dictionary object."""
        return self.validate_model()

    def update(self, other: Self) -> None:
        """Update model content with content of another model."""
        self._model_data.update(other._model_data)


class Model(_MapModel, ABC):
    """Base class for build configuration models."""


class ModelList(_ModelBase, ABC, MutableSequence[_TModelBase]):
    """Base class for lists containing models."""

    def __init__(self, i: Model | Iterable[Model] | None = None, /, *models: Model) -> None:
        super().__init__(
            [i] + list(models) if isinstance(i, Model) else list(i or []) + list(models)
        )

    @classmethod
    def from_list(cls, i: Sequence[Any], /) -> Self:
        """Create model from list object."""
        return cls(i)

    def insert(self, index: int, value: Model) -> None:
        """Insert model before index."""
        self._model_data.insert(index, value)

    def list(self) -> list[Any]:
        """Get model as list object."""
        return self.validate_model()

    @singledispatchmethod
    def __getitem__(self, index) -> Model | Self:  # noqa D105
        raise TypeError(f"list indices must be integers or slices, not {type(index)}")

    @__getitem__.register
    def _(self, index: int) -> Model:
        return self._model_data[index]

    @__getitem__.register
    def _(self, index: slice) -> Self:
        return self.__class__(self._model_data[index])

    def __setitem__(self, index: int | slice, value: Model | Iterable[Model]) -> None:  # noqa D105
        self._model_data[index] = value

    def __delitem__(self, index: int | slice) -> None:  # noqa D105
        del self._model_data[index]

    def __len__(self) -> int:  # noqa D105
        return len(self._model_data)


def _sort_model_list(model_list: ModelList) -> dict[str, ModelList]:
    """Sort a model list by the wrapped model's primary key value."""
    result: dict[str, ModelList] = {}

    primary_key = _glom_primary_key(model_list)
    for model in model_list:
        primary_key_value = getattr(model, primary_key)
        if primary_key_value in result:
            result[primary_key_value].append(model)
        else:
            result[primary_key_value] = model_list.__class__(model)

    return result


class ModelMapping(_MapModel, ABC, MutableMapping[_KT, _TModelBase]):
    """Base class for mappings containing models."""

    def __getitem__(self, key: str, /) -> Any:  # noqa D105
        return self._model_data[key]

    def __setitem__(self, key: str, value: Any, /) -> None:  # noqa D105
        self._model_data[key] = value

    def __delitem__(self, key: str, /) -> None:  # noqa D105
        del self._model_data[key]

    def __iter__(self) -> Iterator[str]:  # noqa D105
        return iter(self._model_data)

    def __len__(self) -> int:  # noqa D105
        return len(self._model_data)


def make_model_builder(
    *models: type[Model] | type[ModelList] | type[ModelMapping],
) -> Callable[[Any], Any]:
    """Make a custom model builder from a provided list of models.

    Args:
        *models: Configuration model to create if json schema successfully matches.

    Returns:
        Callable for building a model from a base Python object.
    """

    def match(obj: Any) -> Model | ModelList | ModelMapping | None:
        for model in models:
            try:
                validate(obj, schema=model.__model_schema__)
                return model(obj)
            except ValidationError:
                pass

        return None

    @subclassdispatch
    def build(obj: Any) -> Any:
        raise TypeError(f"expected list or Mapping, not {type(obj)}")

    @build.register
    def _(obj: Mapping) -> dict[str, Any]:
        result: dict[str, Any] = {}

        for k, v in obj.items():
            if model := match(v):
                built = model
            elif isinstance(v, Mapping | list):
                built = build(v)
            else:
                built = v

            result[k] = built

        return result

    @build.register
    def _(obj: list) -> list[Any]:
        result: list[Any] = []

        for v in obj:
            if model := match(v):
                built = model
            elif isinstance(v, Mapping | list):
                built = build(v)
            else:
                built = v

            result.append(built)

        return result

    return build


def _make_property(name: str, default_factory: Callable[[], Any] | None = None) -> property:
    """Make new property for a data model field.

    Args:
        name: Name of configuration option that a model field maps to.
        default_factory:
            Zero argument callable to invoke when `del` is either called on the model field
            or the model field is accessed but the configuration option is not present in the
            model's internal data map.
    """

    def getter(self: Model) -> Any | None:
        if default_factory and name not in self._model_data:
            self._model_data[name] = default_factory()

        return self._model_data.get(name, None)

    def setter(self: Model, v: Any) -> None:
        self._model_data[name] = v

    def deleter(self: Model) -> None:
        try:
            del self._model_data[name]
            if default_factory:
                self._model_data[name] = default_factory()
        except KeyError:
            pass

    return property(getter, setter, deleter)


def _make_annotation_map(
    model_t: _ModelBase, include_origin: bool = False
) -> dict[str, Annotated]:
    """Make a metadata mapping from a model class.

    Args:
        model_t: Model class to make metadata mapping from.
        include_origin: Include the metadata `origin` field in the annotation mapping.
    """
    annotation_map: dict[str, Annotated] = {}

    for field, annotation in get_type_hints(model_t, include_extras=True).items():
        metadata = _get_metadata(annotation)
        annotation_map[metadata.alias or _format_field(field)] = annotation
        if include_origin and metadata.origin:
            annotation_map[metadata.origin] = annotation

    return annotation_map


def _get_metadata(annotation: type) -> Metadata:
    """Get metadata for a Slurm configuration field.

    Args:
        annotation: Annotation to get metadata of.

    Notes:
        If the field is not annotated with `typing.Annotated`, a
        default `Metadata` object will be returned instead.
    """
    if hasattr(annotation, "__metadata__"):
        return annotation.__metadata__[0]

    return Metadata()


def _get_type(annotation: type) -> tuple[Any, ...]:
    """Get the expected type(s) of a Slurm configuration field.

    Args:
        annotation: Annotation to get the Slurm configuration field type from.

    Notes:
        Some Slurm configuration fields can be multiple types, so a tuple of types is
        returned rather than a single type. For example, the `slurm.conf` configuration
        options `MaxNodes` can be both an integer or the string "unlimited".
    """
    # Check if type annotation is wrapped with `typing.Annotated`.
    if hasattr(annotation, "__origin__"):
        annotation = annotation.__origin__

    # Check if origin is an "Optional". For example, is the annotation `str | None`?
    # Almost all annotations will be `<type> | None` since Slurm configuration
    # fields are nullable.
    if get_origin(annotation) is UnionType:
        return tuple(filterfalse(lambda t: t is NoneType, get_args(annotation)))
    else:
        return (annotation,)


def _get_primary_key(model_t: _ModelBase) -> str | None:
    r"""Get primary key of model.

    Args:
        model_t: Model to get primary key of.

    Returns:
        Model's primary key or `None` if the model has no declared primary key.
    """
    for field, annotation in get_type_hints(model_t, include_extras=True).items():
        metadata = _get_metadata(annotation)
        if metadata.primary:
            return field

    return None


def _glom_primary_key(model_list: ModelList) -> str:
    """Get the primary key of a list of models.

    Args:
        model_list: List of models to get the primary key of.

    Raises:
        ModelError: Raised if not all models in the list have the same primary key.
    """
    primary_keys = Counter([_get_primary_key(model) for model in model_list])
    if len(primary_keys) > 1 or None in primary_keys:
        raise ModelError(
            (
                f"models in {model_list} either have inconsistent primary keys"
                f"or do not have a primary key defined: {primary_keys}"
            )
        )

    return cast(str, list(primary_keys)[0])


def _format_field(field: str) -> str:
    """Format configuration field to be case-insensitive and remove underscores.

    Args:
        field: Configuration field to format.

    Examples:
        `_format_field("slurmctld_host")` -> "slurmctldhost"
    """
    return field.replace("_", "").lower()


@subclassdispatch
def _get_as_builtin_object(obj: Iterable[Any]) -> Iterable[Any]:
    """Recursively expand a complex iterable with nested models.

    Args:
        obj: Complex iterable to expand.
    """
    raise TypeError(f"expected list or Mapping, not {type(obj)}")


@_get_as_builtin_object.register
def _(obj: Mapping) -> dict[str, Any]:
    return {k: _expand_by_type(v) for k, v in obj.items()}


@_get_as_builtin_object.register
def _(obj: list) -> list[Any]:
    return [_expand_by_type(v) for v in obj]


def _expand_by_type(v: Any) -> Any:
    """Expand an object based on its type.

    Args:
        v: Object to expand based on its type.

    Returns:
        An "expanded" builtin object with no nested custom objects.
    """
    match v:
        case Model():
            return _get_as_builtin_object(v.dict())
        case ModelList():
            return _get_as_builtin_object(v.list())
        case Mapping() | list():
            return _get_as_builtin_object(v)
        case _:
            return v


@subclassdispatch
def _parse(model_t, s: str) -> Iterable[Any]:
    """Parse data in Slurm configuration format into models.

    Args:
        model_t: Model to create from parsed Slurm configuration data.
        s: Data in Slurm configuration format.

    Raises:
        TypeError: Raised if the type of `model_t` is not `Model`, `ModelList`, or `ModelMapping`.
        ModelError: Raised if an error occurs when parsing the value of `s`.
    """
    raise TypeError(
        f"expected subclass of type Model, ModelList, or ModelMapping, not {type(model_t)}"
    )


@_parse.register
def _(model_t: Model, s: str) -> dict[str, Any]:
    annotation_map = _make_annotation_map(model_t)

    classification_map: dict[str, list[str]] = {}
    for token, expr in _scan(s, mode=model_t.__model_mode__):
        if token not in annotation_map:
            raise ModelError(
                (
                    f"unrecognized configuration field `{token}`. "
                    f"valid configuration fields are: {list(annotation_map.keys())}"
                )
            )

        classification_map[token] = classification_map.get(token, []) + [expr]

    result: dict[str, Any] = {}
    for token, expr in classification_map.items():
        annotation = annotation_map[token]
        metadata = _get_metadata(annotation)
        tp = _get_type(annotation)

        if metadata.unique and len(expr) > 1:
            raise ModelError(
                (
                    f"expected configuration field `{token}` to be unique, "
                    f"but multiple entries found: {classification_map}"
                )
            )

        if isclass(tp[0]) and issubclass(tp[0], _ModelBase):
            result[metadata.origin or token] = tp[0].from_str("\n".join(expr))
            continue

        if metadata.unique:
            result[metadata.origin or token] = metadata.parse(expr[0])
        else:
            result[metadata.origin or token] = [metadata.parse(e) for e in expr]

    return result


@_parse.register
def _(model_t: ModelList, s: str) -> list[Any]:
    result: list[Any] = []
    origin = get_original_bases(cast(type, model_t))[0]
    wrapped_model_t = get_args(origin)[0]
    for _, expr in _scan(s, model_t.__model_mode__):
        result.append(wrapped_model_t.from_str(expr))

    return result


@_parse.register
def _(model_t: ModelMapping, s: str) -> dict[str, Any]:
    result: dict[str, Any] = {}
    origin = get_original_bases(cast(type, model_t))[0]
    wrapped_model_t = get_args(origin)[1]
    primary_key = _get_primary_key(wrapped_model_t)

    if primary_key:
        for _, expr in _scan(s, model_t.__model_mode__):
            model = wrapped_model_t.from_str(expr)
            result.update({getattr(model, primary_key): model})
    elif primary_key is None and issubclass(wrapped_model_t, ModelList):
        model = wrapped_model_t.from_str(s)
        result.update(_sort_model_list(model))
    else:
        raise ModelError(
            (
                f"model {wrapped_model_t} has no primary key. "
                f"{wrapped_model_t} cannot be used with a model mapping"
            )
        )

    return result


def _clean(s: str) -> list[str]:
    """Clean configuration line(s) before further processing.

    Args:
        s: Configuration line(s) to clean.

    Warnings:
        This function will strip all comments from the configuration line(s),
        but will not preserve them. Original comments will not be present in
        the marshalled configuration data output.
    """
    result: list[str] = []

    for line in s.splitlines():
        cleaned = line.split("#", maxsplit=1)[0].strip()
        if cleaned == "":  # Line was a comment line or is empty.
            continue

        result.append(cleaned)

    return result


def _scan(s: str, /, mode: Mode) -> Iterable[tuple[str, str]]:
    """Scan Slurm configuration data to identify key-value pairs.

    Args:
        s: Configuration data to scan for key-value pairs.
        mode: Mode to scan file in.
    """
    match mode:
        case mode.ONELINE:
            expr = _scan_line(c[0]) if (c := _clean(s)) else []

        case mode.STANZA:
            expr = _clean(s)

        case _:
            raise TypeError(f"mode must be either ONELINE or STANZA, not {type(mode)}")

    for e in expr:
        token = "".join(map(str, takewhile(lambda c: c != "=" and c != " ", e))).lower()
        yield token, e


def _scan_line(s: str) -> list[str]:
    """Scan a Slurm configuration line to identify key-value pairs.

    Args:
        s: Configuration line to scan for key-value pairs.

    Notes:
        This method is used rather than `shlex.split(...)` or a regular expression
        because this method offers a ~429% performance increase over `shlex.split(...)`
        for identifying key-value pairs in a Slurm configuration line.
    """
    pairs: list[str] = []
    quote = None
    current_pair = ""

    for c in s:
        if c in ['"', "'"]:
            if c == quote:
                quote = None
                current_pair += c
            elif quote is None:
                quote = c
                current_pair += c
            else:
                current_pair += c
        elif c.isspace() and quote is None:
            if current_pair:
                pairs.append(current_pair)
                current_pair = ""
        else:
            current_pair += c

    if current_pair:
        pairs.append(current_pair)

    return pairs


@subclassdispatch
def _marshal(model) -> str:
    """Marshal a model into Slurm configuration format.

    Args:
        model: Configuration model to marshal into Slurm configuration format.
    """
    raise TypeError(
        f"expected subclass of type Model, ModelList, or ModelMapping, not {type(model)}"
    )


@_marshal.register
def _(model: Model) -> str:
    annotation_map = _make_annotation_map(model, include_origin=True)
    result: deque[str] = deque()

    # Directly access `_model_data` here because `_expand(...)` does not
    # preserve the type information required for subclass dispatching.
    for k, v in model._model_data.items():  # noqa SLF001
        if isinstance(v, Model | ModelList | ModelMapping):
            result.append(str(v))
            continue

        metadata = _get_metadata(annotation_map[k])
        expr = metadata.marshal(metadata.alias or k, v)
        if metadata.primary:
            result.appendleft(expr)
        else:
            result.append(expr)

    return _format_output(result, mode=model.__model_mode__)


@_marshal.register
def _(model: ModelList) -> str:
    result: list[str] = []
    for wrapped_model in model:
        result.append(str(wrapped_model))

    return _format_output(result, mode=model.__model_mode__)


@_marshal.register
def _(model: ModelMapping) -> str:
    result: list[str] = []
    for wrapped_model in model.values():
        result.append(str(wrapped_model))

    return _format_output(result, mode=model.__model_mode__)


def _format_output(i: Iterable[str], mode: Mode) -> str:
    """Format `_marshal` output based on model mode.

    Args:
        i: Iterable to format based on given mode.
        mode: Mode to format output in.

    Raises:
        TypeError:
            Raised if invalid marshaling mode is provided.
            Valid modes are `ONELINE` and `STANZA`.
    """
    match mode:
        case Mode.ONELINE:
            return " ".join(i)
        case Mode.STANZA:
            return "\n".join(i)
        case _:
            raise TypeError(f"mode must be either ONELINE or STANZA, not {type(mode)}")
