from collections.abc import ItemsView, Iterator, KeysView, MutableMapping, Set
from types import GenericAlias, UnionType
from typing import (
    ClassVar,
    Never,
    Protocol,
    TypeAliasType,
    cast,
    get_args,
    get_origin,
    overload,
)


class TypedYamlDict[K: object, V: object](Protocol):
    def __getitem__(self, key: str | K, /) -> V: ...
    # def __setitem__(self, key: str, value: V, /) -> V: ...
    def __delitem__(self, key: Never | K, /) -> None: ...
    def __contains__(self, key: K, /) -> bool: ...
    def __iter__(self) -> Iterator[K]: ...
    def __len__(self) -> int: ...
    def keys(self) -> KeysView[K]: ...
    def items(self) -> ItemsView[K, V]: ...
    def pop(self, key: Never | K, /) -> V: ...

    # def popitem(self) -> tuple[K, V]: ...

    # def clear(self) -> None: ...

    __required_keys__: ClassVar[frozenset[str]]
    __optional_keys__: ClassVar[frozenset[str]]


# class Test(TypedDict):
#     var: str
#
#
# x = Test(var="test")
#
#
# def is_typed_dict_test(obj: TypedYamlDict[object, object]) -> None:
#     print(obj)
#     pass
#
#
# is_typed_dict_test(x)

type T_Primitive = None | bool | int | str

type T_PrimIters = tuple[T_Prim, ...] | list[T_Prim] | Set[T_Prim] | Iterator[T_Prim]
type T_PrimDict = MutableMapping[T_Primitive, T_Prim]
type T_Prim = T_Primitive | T_PrimIters | T_PrimDict

type T_YamlIters = tuple[T_Yaml, ...] | list[T_Yaml] | Set[T_Yaml] | Iterator[T_Yaml]
type T_YamlDict = MutableMapping[str, T_Yaml]
type T_YamlRW = T_YamlIters | T_YamlDict
type T_Yaml = T_Primitive | T_YamlRW


type T_YamlPostDict = TypedYamlDict[str, T_YamlPost]
type T_YamlPostRes = tuple[T_YamlPost, ...] | T_YamlPostDict
type T_YamlPost = T_Primitive | T_YamlPostRes


def get_union_types(annotations: UnionType) -> Iterator[type]:
    for annotation in get_args(annotations):  # pyright: ignore[reportAny]
        if isinstance(annotation, TypeAliasType):
            annotation = annotation.__value__  # pyright: ignore[reportAny]
        if isinstance(annotation, UnionType):
            yield from get_union_types(annotation)
            continue
        yield get_types(annotation)  # pyright: ignore[reportAny]


@overload
def get_types(annotation: UnionType) -> tuple[type]:
    pass


@overload
def get_types(annotation: GenericAlias) -> type:
    pass


@overload
def get_types(annotation: TypeAliasType) -> type | tuple[type, ...]:
    pass


def get_types(
    annotation: TypeAliasType | GenericAlias | UnionType,
) -> type | tuple[type, ...]:
    if isinstance(annotation, TypeAliasType):
        annotation = annotation.__value__  # pyright: ignore[reportAny]
    if isinstance(annotation, GenericAlias):
        return get_origin(annotation)
    if isinstance(annotation, UnionType):
        return tuple(get_union_types(annotation))
    return cast(type, annotation)  # pyright: ignore[reportInvalidCast]