"""Make `hypothesis` resolve types to the right strategies."""

from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from typing import TypeAlias, TypeVar

from hypothesis import strategies as st
from hypothesis.strategies._internal import types  # noqa: PLC2701

Example_co = TypeVar('Example_co', covariant=True)

StrategyFactory: TypeAlias = (
    st.SearchStrategy[Example_co]
    | Callable[[type[Example_co]], st.SearchStrategy[Example_co]]
)


@contextmanager
def strategies_for_types(
    mapping: Mapping[type[object], StrategyFactory],
) -> Iterator[None]:
    """
    Temporarily register strategies with `hypothesis`.

    Within this context, `hypothesis` will generate data for `MyType`
    using `mapping[MyType]`, if available. Otherwise, it will continue to
    use the globally registered strategy for `MyType`.

    NOTE: This manually adds and removes strategies from an internal data
    structure of `hypothesis`: `types._global_type_lookup`. This is a global
    variable used for practically every example generated by `hypothesis`, so
    we can easily have unintentional side-effects. We have to be very careful
    when modifying it.
    """
    previous_strategies: dict[type[object], StrategyFactory | None] = {}
    for type_, strategy in mapping.items():
        previous_strategies[type_] = look_up_strategy(type_)
        st.register_type_strategy(type_, strategy)

    try:
        yield
    finally:
        for type_, previous_strategy in previous_strategies.items():
            if previous_strategy is None:
                _remove_strategy(type_)
            else:
                st.register_type_strategy(type_, previous_strategy)


def look_up_strategy(
    type_: type[Example_co],
) -> StrategyFactory[Example_co] | None:
    """Return the strategy used by `hypothesis`."""
    return types._global_type_lookup.get(type_)  # noqa: SLF001


def _remove_strategy(
    type_: type[object],
) -> None:
    """Remove the strategy registered for `type_`."""
    types._global_type_lookup.pop(type_)  # noqa: SLF001
    _clean_caches()


def apply_strategy(
    strategy: StrategyFactory[Example_co], type_: type[Example_co]
) -> StrategyFactory[Example_co]:
    """Apply `strategy` to `type_`."""
    if isinstance(strategy, st.SearchStrategy):
        return strategy
    return strategy(type_)


def _clean_caches() -> None:
    st.from_type.__clear_cache()  # type: ignore[attr-defined]  # noqa: SLF001
