"""Stub file for hass_dict. Provide overload for type checking."""
# ruff: noqa: PYI021  # Allow docstrings

from typing import Any, Generic, TypeVar, assert_type, overload

__all__ = [
    "HassDict",
    "HassEntryKey",
    "HassKey",
]

_T = TypeVar("_T")  # needs to be invariant

class _Key(Generic[_T]):
    """Base class for Hass key types. At runtime delegated to str."""

    def __init__(self, value: str, /) -> None: ...
    def __len__(self) -> int: ...
    def __hash__(self) -> int: ...
    def __eq__(self, other: object) -> bool: ...
    def __getitem__(self, index: int) -> str: ...

class HassEntryKey(_Key[_T]):
    """Key type for integrations with config entries."""

class HassKey(_Key[_T]):
    """Generic Hass key type."""

class HassDict(dict[_Key[Any] | str, Any]):
    """Custom dict type to provide better value type hints for Hass key types."""

    @overload  # type: ignore[override]
    def __getitem__[_S](self, key: HassEntryKey[_S], /) -> dict[str, _S]: ...
    @overload
    def __getitem__[_S](self, key: HassKey[_S], /) -> _S: ...
    @overload
    def __getitem__(self, key: str, /) -> Any: ...

    # ------
    @overload  # type: ignore[override]
    def __setitem__[_S](
        self, key: HassEntryKey[_S], value: dict[str, _S], /
    ) -> None: ...
    @overload
    def __setitem__[_S](self, key: HassKey[_S], value: _S, /) -> None: ...
    @overload
    def __setitem__(self, key: str, value: Any, /) -> None: ...

    # ------
    @overload  # type: ignore[override]
    def setdefault[_S](
        self, key: HassEntryKey[_S], default: dict[str, _S], /
    ) -> dict[str, _S]: ...
    @overload
    def setdefault[_S](self, key: HassKey[_S], default: _S, /) -> _S: ...
    @overload
    def setdefault(self, key: str, default: None = None, /) -> Any | None: ...
    @overload
    def setdefault(self, key: str, default: Any, /) -> Any: ...

    # ------
    @overload  # type: ignore[override]
    def get[_S](self, key: HassEntryKey[_S], /) -> dict[str, _S] | None: ...
    @overload
    def get[_S, _U](
        self, key: HassEntryKey[_S], default: _U, /
    ) -> dict[str, _S] | _U: ...
    @overload
    def get[_S](self, key: HassKey[_S], /) -> _S | None: ...
    @overload
    def get[_S, _U](self, key: HassKey[_S], default: _U, /) -> _S | _U: ...
    @overload
    def get(self, key: str, /) -> Any | None: ...
    @overload
    def get(self, key: str, default: Any, /) -> Any: ...

    # ------
    @overload  # type: ignore[override]
    def pop[_S](self, key: HassEntryKey[_S], /) -> dict[str, _S]: ...
    @overload
    def pop[_S](
        self, key: HassEntryKey[_S], default: dict[str, _S], /
    ) -> dict[str, _S]: ...
    @overload
    def pop[_S, _U](
        self, key: HassEntryKey[_S], default: _U, /
    ) -> dict[str, _S] | _U: ...
    @overload
    def pop[_S](self, key: HassKey[_S], /) -> _S: ...
    @overload
    def pop[_S](self, key: HassKey[_S], default: _S, /) -> _S: ...
    @overload
    def pop[_S, _U](self, key: HassKey[_S], default: _U, /) -> _S | _U: ...
    @overload
    def pop(self, key: str, /) -> Any: ...
    @overload
    def pop[_U](self, key: str, default: _U, /) -> Any | _U: ...

def _test_hass_dict_typing() -> None:  # noqa: PYI048
    """Test HassDict overloads work as intended.

    This is tested during the mypy run. Do not move it to 'tests'!
    """
    d = HassDict()
    entry_key = HassEntryKey[int]("entry_key")
    key = HassKey[int]("key")
    key2 = HassKey[dict[int, bool]]("key2")
    key3 = HassKey[set[str]]("key3")
    other_key = "domain"

    # __getitem__
    assert_type(d[entry_key], dict[str, int])
    assert_type(d[entry_key]["entry_id"], int)
    assert_type(d[key], int)
    assert_type(d[key2], dict[int, bool])

    # __setitem__
    d[entry_key] = {}
    d[entry_key] = 2  # type: ignore[call-overload]
    d[entry_key]["entry_id"] = 2
    d[entry_key]["entry_id"] = "Hello World"  # type: ignore[assignment]
    d[key] = 2
    d[key] = "Hello World"  # type: ignore[misc]
    d[key] = {}  # type: ignore[misc]
    d[key2] = {}
    d[key2] = 2  # type: ignore[misc]
    d[key3] = set()
    d[key3] = 2  # type: ignore[misc]
    d[other_key] = 2
    d[other_key] = "Hello World"

    # get
    assert_type(d.get(entry_key), dict[str, int] | None)
    assert_type(d.get(entry_key, True), dict[str, int] | bool)
    assert_type(d.get(key), int | None)
    assert_type(d.get(key, True), int | bool)
    assert_type(d.get(key2), dict[int, bool] | None)
    assert_type(d.get(key2, {}), dict[int, bool])
    assert_type(d.get(key3), set[str] | None)
    assert_type(d.get(key3, set()), set[str])
    assert_type(d.get(other_key), Any | None)
    assert_type(d.get(other_key, True), Any)
    assert_type(d.get(other_key, {})["id"], Any)

    # setdefault
    assert_type(d.setdefault(entry_key, {}), dict[str, int])
    assert_type(d.setdefault(entry_key, {})["entry_id"], int)
    assert_type(d.setdefault(key, 2), int)
    assert_type(d.setdefault(key2, {}), dict[int, bool])
    assert_type(d.setdefault(key2, {})[2], bool)
    assert_type(d.setdefault(key3, set()), set[str])
    assert_type(d.setdefault(other_key, 2), Any)
    assert_type(d.setdefault(other_key), Any | None)
    d.setdefault(entry_key, {})["entry_id"] = 2
    d.setdefault(entry_key, {})["entry_id"] = "Hello World"  # type: ignore[assignment]
    d.setdefault(key, 2)
    d.setdefault(key, "Error")  # type: ignore[misc]
    d.setdefault(key2, {})[2] = True
    d.setdefault(key2, {})[2] = "Error"  # type: ignore[assignment]
    d.setdefault(key3, set()).add("Hello World")
    d.setdefault(key3, set()).add(2)  # type: ignore[arg-type]
    d.setdefault(other_key, {})["id"] = 2
    d.setdefault(other_key, {})["id"] = "Hello World"
    d.setdefault(entry_key)  # type: ignore[call-overload]
    d.setdefault(key)  # type: ignore[call-overload]
    d.setdefault(key2)  # type: ignore[call-overload]

    # pop
    assert_type(d.pop(entry_key), dict[str, int])
    assert_type(d.pop(entry_key, {}), dict[str, int])
    assert_type(d.pop(entry_key, 2), dict[str, int] | int)
    assert_type(d.pop(key), int)
    assert_type(d.pop(key, 2), int)
    assert_type(d.pop(key, "Hello World"), int | str)
    assert_type(d.pop(key2), dict[int, bool])
    assert_type(d.pop(key2, {}), dict[int, bool])
    assert_type(d.pop(key2, 2), dict[int, bool] | int)
    assert_type(d.pop(key3), set[str])
    assert_type(d.pop(key3, set()), set[str])
    assert_type(d.pop(other_key), Any)
    assert_type(d.pop(other_key, True), Any | bool)