from __future__ import annotations

import sys
from abc import ABC, ABCMeta
from collections.abc import Callable
from typing import get_type_hints

__all__ = ["AnnotatedABC", "AnnotatedABCMeta"]


def _normalize_annotations(fn: Callable[..., object]) -> dict[str, object]:
    module_name = getattr(fn, "__module__", None)
    module = sys.modules.get(module_name) if module_name is not None else None
    globalns = vars(module) if module is not None else {}
    return get_type_hints(fn, globalns=globalns, localns=globalns)


def _has_fn(
    impl: type,
    abc: type,
    name: str,
    expected: dict[str, object],
) -> bool:
    skip = set(abc.__mro__)
    for base in impl.__mro__:
        if base in skip:
            continue
        for key, value in base.__dict__.items():
            if key == name and _normalize_annotations(value) == expected:
                return True
    return False


class AnnotatedABCMeta(ABCMeta):
    def register(cls, subclass: type) -> type:
        cls._validate_implementation(subclass)
        return super().register(subclass)

    def _validate_implementation(cls, impl: type) -> None:
        for name, expected in cls._abstract_annotations().items():
            if not _has_fn(impl, cls, name, expected):
                msg = (
                    f"{impl!r} does not fully implement {cls.__name__!r}: "
                    f"missing or mismatched {name!r}"
                )
                raise TypeError(msg)

    def _abstract_annotations(cls) -> dict[str, dict[str, object]]:
        methods: dict[str, dict[str, object]] = {}
        for base in cls.__mro__:
            if not isinstance(base, AnnotatedABCMeta):
                continue
            for name, value in base.__dict__.items():
                if getattr(value, "__isabstractmethod__", False):
                    methods[name] = _normalize_annotations(value)
        return methods


class AnnotatedABC(ABC, metaclass=AnnotatedABCMeta):
    def __init_subclass__(cls, **kwargs: object) -> None:
        super().__init_subclass__(**kwargs)
        for base in cls.__bases__:
            if isinstance(base, AnnotatedABCMeta):
                base._validate_implementation(cls)
