from collections.abc import MutableMapping
from typing import Any, Generic, Self, TypeVar

from pydantic import BaseModel as PydanticBaseModel
from pydantic import ConfigDict
from pydantic import RootModel as PydanticRootModel

from .serialization import PydanticTomlMixin, PydanticYamlMixin

T = TypeVar("T")


class BaseModel(
    PydanticBaseModel,
    PydanticTomlMixin,
    PydanticYamlMixin,
):
    pass


class RootModel(
    PydanticRootModel[T],
    PydanticTomlMixin,
    PydanticYamlMixin,
    Generic[T],
):
    pass


_immutable_model_config = ConfigDict(
    frozen=True,
    validate_assignment=True,
)


class _ImmutableMixin:
    """
    A base model that is immutable.
    """

    def __init_subclass__(cls, **kwargs):
        """
        Add the immutable model config to the rest of this class' model config.
        """
        cls.model_config = cls.model_config | _immutable_model_config
        super().__init_subclass__(**kwargs)


class ImmutableBaseModel(BaseModel, _ImmutableMixin):
    def model_update(self, **kwargs: Any) -> Self:
        """
        Returns a new copy of the model with the given kwargs updated, and
        validates the assignment.
        Contrary to common belief, setting revalidate_instances="always" and
        validate_assignment=True is not enough, because model_copy bypasses all
        of that.
        See https://github.com/pydantic/pydantic/issues/418
        """
        assert isinstance(self, BaseModel)
        updated = self.model_copy(update=kwargs)  # type: ignore
        self.__class__.model_validate(updated.__dict__)  # validate the updated model

        # Clear cached properties that may have been copied from the original instance
        # model_copy copies __dict__ which includes functools.cached_property values
        model_field_names = set(self.__class__.model_fields.keys())
        cached_keys = [k for k in updated.__dict__ if k not in model_field_names]
        updated_dict = updated.__dict__
        assert isinstance(updated_dict, MutableMapping)
        for key in cached_keys:
            del updated_dict[key]

        return updated

    def _model_mutate(self, **kwargs: Any):
        """
        Mutates the model in place despite its immutability.
        This is obviously dangerous, so the method is protected.
        To make it slightly safer, we validate the assignment via model_update.
        """
        updated = self.model_update(**kwargs)
        assert isinstance(self.__dict__, MutableMapping)
        self.__dict__.update(updated.__dict__)


T = TypeVar("T")


class ImmutableRootModel(RootModel[T], _ImmutableMixin, Generic[T]):
    pass


class MutableBaseModel(BaseModel):
    pass


class MutableRootModel(RootModel[T], Generic[T]):
    pass


__all__ = [
    "ImmutableBaseModel",
    "ImmutableRootModel",
    "MutableBaseModel",
    "MutableRootModel",
]
