from abc import ABC
from collections import defaultdict
from collections.abc import Callable
from typing import Self, TypeVar, cast, get_args

__all__ = ("FrozenDefaultDict", "FrozenDict", "Immutable", "is_immutable")

T = TypeVar("T")
K = TypeVar("K")


type Immutable = int | float | tuple | str | FrozenDict | FrozenDefaultDict
_IMMUTABLE_TYPES = get_args(Immutable)


def is_immutable(ob: object) -> bool:
    return isinstance(ob, _IMMUTABLE_TYPES)


def freeze(ob: object) -> object:
    if is_immutable(ob):
        return ob
    if isinstance(ob, list | set):
        return tuple(ob)
    if isinstance(ob, dict):
        return FrozenDict(ob)
    if isinstance(ob, defaultdict):
        return FrozenDefaultDict.from_defaultdict(ob)
    raise NotImplementedError
    # TODO: overload to convert any object into its immutable counterpart (if not already)


class _FrozenDictBase(ABC):
    @property
    def _immutable_error(self) -> TypeError:
        return TypeError(f"{self.__class__.__name__} is immutable.")

    def update(self, _) -> None:
        raise self._immutable_error

    def setdefault(self, _, __) -> None:
        raise self._immutable_error

    def pop(self) -> None:
        raise self._immutable_error

    def popitem(self, _) -> None:
        raise self._immutable_error

    def __setitem__(self, _, __) -> None:
        raise self._immutable_error

    def __delitem__(self, _) -> None:
        raise self._immutable_error


class FrozenDict(_FrozenDictBase, dict[K, T]):
    ...  # type: ignore
    # see https://pypi.org/project/frozendict/, but make stricter


class FrozenDefaultDict[K, T](_FrozenDictBase, defaultdict[K, T]):  # type: ignore
    def __init__(self, default_factory: Callable[[], T], dictionary: dict[K, T]):
        super().__init__(default_factory, dictionary)

    @classmethod
    def from_defaultdict(cls, dd: defaultdict[K, T]) -> Self:
        _factory = cast(Callable[[], T], dd.default_factory)
        return cls(_factory, dict(dd.items()))

    def __repr__(self) -> str:
        default = self.default_factory
        result = default.__name__ if callable(default) else repr(default)
        return f"{self.__class__.__name__}({dict(self)}, default={result})"

    def __getitem__(self, key: K) -> T:
        if key in self:
            return super().__getitem__(key)
        return cast(Callable, self.default_factory)()
