import asyncio
import inspect
from collections import UserDict
from collections.abc import Awaitable, Callable, Iterator
from contextlib import suppress
from functools import wraps
from typing import Any, Literal, ParamSpec, Self, TypeVar, overload
from unittest.mock import MagicMock, create_autospec, seal

import pytest
from fastapi import FastAPI
from fastapi.dependencies.utils import is_coroutine_callable

_T = TypeVar("_T")
_P = ParamSpec("_P")

_DepType = Callable[_P, _T]


class Overrider(UserDict):
    """Set dependency overrides and clean the up after using.
    To be used as a pytest fixture."""

    def __init__(
        self,
        app: FastAPI,
    ) -> None:
        self._app = app

    @overload
    def __call__(self, key: _DepType, override: _DepType) -> _DepType:
        """Override a dependency with the given function.
        Returns the function"""
        ...

    @overload
    def __call__(self, key: _DepType, override: _T) -> _T:
        """Override a dependeny with a function returning the given value.
        Returns the value"""
        ...

    @overload
    def __call__(self, key: _DepType) -> MagicMock:
        """Override a dependnecy with a mock value.
        Returns the mock value"""
        ...

    def __call__(self, *args, **kwargs) -> _DepType | MagicMock | object:
        """Override a dependency either with a function, a value or a mock."""
        match args:
            case [key] if isinstance(key, Callable) and (
                list(kwargs.keys()) == ["strict"] or len(kwargs) == 0
            ):
                return self.mock(key)
            case [key, override] if isinstance(key, Callable) and isinstance(
                override, Callable
            ):
                return self.function(key, override)
            case [key, override] if isinstance(key, Callable):
                return self.value(key, override)
            case _:
                raise NotImplementedError

    def function(self, key: _DepType, override: _DepType) -> _DepType:
        """Override a dependency with the given function.
        Returns the function"""
        self[key] = override
        return override

    def value(self, key: _DepType, override: _T) -> _T:
        """Override a dependeny with a function returning the given value.
        Returns the value"""

        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:  # noqa: ARG001
            return override

        self[key] = wraps(key)(wrapper)
        return override

    def mock(self, key: _DepType) -> MagicMock:
        """Override a dependnecy with a mock value.
        Returns a mock function that returns a mock value"""
        value_name = f"mock value for {key.__name__}"
        function_name = f"mock function for {key.__name__}"
        return_type = self._get_return_type(key)
        return_value = create_autospec(
            return_type,
            instance=True,
            unsafe=False,
            name=value_name,
        )

        mock_func = MagicMock(unsafe=False, name=function_name)
        mock_func.__signature__ = inspect.signature(key)
        mock_func.return_value = return_value
        seal(mock_func)
        self[key] = mock_func
        return mock_func

    def spy(self, key: _DepType) -> MagicMock:
        """Replace a dependency with a spy wrapper.
        Returns the spy"""
        spy_name = f"Spy for {key.__name__}"

        def wrapper(*args, **kwargs) -> Any:  # noqa: ANN002,ANN003,ANN401
            if is_coroutine_callable(key):
                return asyncio.run(key(*args, **kwargs))
            return key(*args, **kwargs)

        spy = MagicMock(wraps=key, unsafe=False, name=spy_name)
        spy.__signature__ = inspect.signature(key)
        spy.side_effect = wrapper
        seal(spy)
        self[key] = spy
        return spy

    with suppress(ModuleNotFoundError):
        import unifactory

        def some(
            self,
            key: Callable[_P, _T] | Callable[_P, Awaitable[_T]],
            **kwargs: Any,  # noqa: ANN401
        ) -> _T:
            """Override a dependency with a value generated by Polyfactory.
            Additional keyword args are forwarded to the factory's build() method.
            Returns the generated value"""
            override = self.unifactory.build(self._get_return_type(key), **kwargs)
            return self.value(key, override)

        def batch(
            self,
            key: Callable[_P, _T] | Callable[_P, Awaitable[_T]],
            size: int,
            **kwargs: Any,  # noqa: ANN401
        ) -> Iterator[_T]:
            """Batch override a dependency with a given number of values generated by
            Polyfactory.
            Additional keyword args are forwarded to the factory's build() method.
            Returns an Iterator of generated values"""
            overrides = self.unifactory.batch(
                self._get_return_type(key), size, **kwargs
            )
            for override in overrides:
                yield self.value(key, override)

        def cover(
            self,
            key: Callable[_P, _T] | Callable[_P, Awaitable[_T]],
            **kwargs: Any,  # noqa: ANN401
        ) -> Iterator[_T]:
            """Override a dependency with values until full coverage has been reached.
            Additional keyword args are forwarded to the factory's coverage() method.
            Returns an Iterator of generated values"""
            overrides = self.unifactory.coverage(self._get_return_type(key), **kwargs)
            for override in overrides:
                yield self.value(key, override)

    def _get_return_type(
        self, key: Callable[_P, _T] | Callable[_P, Awaitable[_T]]
    ) -> type[_T]:
        return inspect.get_annotations(key)["return"]

    def __enter__(self: Self) -> Self:
        self._restore_overrides = self._app.dependency_overrides
        self._app.dependency_overrides = self._restore_overrides.copy()
        self.data = self._app.dependency_overrides
        return self

    def __exit__(self, *_: object) -> None:
        self._app.dependency_overrides = self._restore_overrides


Scope = (
    Literal["session", "package", "module", "class", "function"]
    | Callable[
        [str, pytest.Config],
        Literal["session", "package", "module", "class", "function"],
    ]
)


def register_fixture(
    app: FastAPI | None = None,
    *,
    name: str = "override",
    autouse: bool = False,
    scope: Scope = "function",
) -> Callable[[], Iterator[Overrider]]:
    if app:

        def inner() -> Iterator[Overrider]:
            with Overrider(app) as overrider:
                yield overrider
    else:

        def inner(app: FastAPI) -> Iterator[Overrider]:
            with Overrider(app) as overrider:
                yield overrider

    return pytest.fixture(name=name, autouse=autouse, scope=scope)(inner)
