from __future__ import annotations

import asyncio
import functools
import inspect
import logging
import threading
from collections.abc import (
    AsyncGenerator,
    Awaitable,
    Callable,
    Generator,
    Iterable,
    Iterator,
)
from contextlib import asynccontextmanager, contextmanager
from dataclasses import asdict, dataclass
from typing import Annotated, Any, NamedTuple, ParamSpec, TypeVar, cast, get_origin

from picodi._internal import NullAwaitable
from picodi._scopes import (
    GlobalScope,
    NullScope,
    ParentCallScope,
    Scope,
    SingletonScope,
)

try:
    import fastapi.params
except ImportError:
    fastapi = None  # type: ignore[assignment] # pragma: no cover


logger = logging.getLogger("picodi")


DependencyCallable = Callable[..., Any]
T = TypeVar("T")
P = ParamSpec("P")
TC = TypeVar("TC", bound=Callable)

unset = object()


class RegistryStorage:
    def __init__(self) -> None:
        self.deps: dict[DependencyCallable, Provider] = {}
        self.overrides: dict[DependencyCallable, DependencyCallable] = {}
        self.lock = threading.RLock()

    def __iter__(self) -> Iterator[Provider]:
        return iter(self.deps.values())


class InternalRegistry:
    def __init__(self, storage: RegistryStorage) -> None:
        self._storage = storage

    def add(
        self,
        dependency: DependencyCallable,
        scope_class: type[Scope] = NullScope,
        override_scope: bool = False,
    ) -> None:
        """
        Add a dependency to the registry.
        """
        with self._storage.lock:
            if dependency in self._storage.deps:
                provider = self._storage.deps[dependency]
                to_replace = provider.replace(
                    scope_class=(scope_class if override_scope else None)
                )
                if to_replace != provider:
                    self._storage.deps[dependency] = to_replace
            else:
                self._storage.deps[dependency] = Provider.from_dependency(
                    dependency=dependency,
                    scope_class=scope_class,
                )

    def get(self, dependency: DependencyCallable) -> Provider:
        with self._storage.lock:
            if self._storage.overrides.get(dependency):
                dependency = self._storage.overrides[dependency]
            return self._storage.deps[dependency]

    def filter(self, predicate: Callable[[Provider], bool]) -> Iterable[Provider]:
        return filter(predicate, self._storage)


class Registry:
    def __init__(
        self, storage: RegistryStorage, internal_registry: InternalRegistry
    ) -> None:
        self._storage = storage
        self._internal_registry = internal_registry

    def override(
        self,
        dependency: DependencyCallable,
        new_dependency: DependencyCallable | None | object = unset,
    ) -> Callable[[DependencyCallable], DependencyCallable]:
        """
        Override a dependency with a new one. It can be used as a decorator,
        as a context manager or as a regular method call. New dependency will be
        added to the registry.
        Examples:
        ```
        @registry.override(get_settings)
        def real_settings():
            return {"real": "settings"}

        with registry.override(get_settings, real_settings):
            ...

        registry.override(get_settings, real_settings)
        registry.override(get_settings, None)  # clear override
        """

        def decorator(override_to: DependencyCallable) -> DependencyCallable:
            self._internal_registry.add(override_to)
            if dependency is override_to:
                raise ValueError("Cannot override a dependency with itself")
            self._storage.overrides[dependency] = override_to
            return override_to

        if new_dependency is unset:
            return decorator

        with self._storage.lock:
            call_dependency = self._storage.overrides.get(dependency)
            if callable(new_dependency):
                decorator(new_dependency)
            else:
                self._storage.overrides.pop(dependency, None)

        @contextmanager
        def manage_context() -> Generator[None, None, None]:
            try:
                yield
            finally:
                self.override(dependency, call_dependency)

        return manage_context()

    def clear(self) -> None:
        """
        Clear the registry. It will remove all dependencies and overrides.
        This method will not close any dependencies. So you need to manually call
        `shutdown_dependencies` before this method.
        """
        with self._storage.lock:
            self._storage.deps.clear()
            self._storage.overrides.clear()

    def clear_overrides(self) -> None:
        """
        Clear all overrides. It will remove all overrides, but keep the dependencies.
        """
        with self._storage.lock:
            self._storage.overrides.clear()


_registry_storage = RegistryStorage()
_internal_registry = InternalRegistry(_registry_storage)
registry = Registry(_registry_storage, _internal_registry)
_scopes: dict[type[Scope], Scope] = {
    NullScope: NullScope(),
    SingletonScope: SingletonScope(),
    ParentCallScope: ParentCallScope(),
}
_lock = threading.RLock()


def Provide(dependency: DependencyCallable, /) -> Any:  # noqa: N802
    """
    Declare a provider.
    It takes a single "dependency" callable (like a function).
    Don't call it directly, picodi will call it for you.
    DependencyCallable can be a regular function or a generator with one yield.
    If the dependency is a generator, it will be used as a context manager.
    Any generator that is valid for `contextlib.contextmanager`
    can be used as a dependency.

    Example:
    ```
    from picodi import Provide, inject

    def get_db():
        yield "db connection"
        print("closing db connection")

    @inject
    def my_service(db: str = Provide(get_db)):
        assert db == "db connection"
    ```
    """
    return Dependency(dependency)


def inject(fn: Callable[P, T]) -> Callable[P, T]:
    """
    Decorator to inject dependencies into a function.
    Use it in combination with `Provide` to declare dependencies.
    Should be placed first in the decorator chain (on bottom).

    Example:
    ```
    from picodi import inject, Provide

    @inject
    def my_service(db=Provide(some_dependency_func)):
        ...
    ```
    """
    signature = inspect.signature(fn)
    dependant = _build_depend_tree(Dependency(fn))

    if inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn):

        @functools.wraps(fn)
        async def fun_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            gen = _wrapper_helper(
                dependant, signature, is_async=True, args=args, kwargs=kwargs
            )
            value, action = next(gen)
            result = None
            exceptions = []
            while True:
                if inspect.iscoroutine(value):
                    try:
                        value = await value
                    except Exception as e:  # noqa: PIE786
                        exceptions.append(e)

                if action == "result":
                    result = value
                try:
                    value, action = gen.send(value)
                except StopIteration:
                    break
            if exceptions:
                # TODO use `ExceptionGroup` after dropping 3.10 support
                raise exceptions[0]
            return cast("T", result)

        wrapper = fun_wrapper

        if inspect.isasyncgenfunction(fn):

            @functools.wraps(fn)
            async def gen_wrapper(
                *args: P.args, **kwargs: P.kwargs
            ) -> AsyncGenerator[T, None]:
                result = await fun_wrapper(*args, **kwargs)
                async for value in result:  # type: ignore[attr-defined]
                    try:
                        yield value
                    except Exception as e:  # noqa: PIE786
                        try:
                            await result.athrow(e)  # type: ignore[attr-defined]
                        except StopAsyncIteration:
                            break

            wrapper = gen_wrapper  # type: ignore[assignment]

    else:

        @functools.wraps(fn)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            gen = _wrapper_helper(
                dependant, signature, is_async=False, args=args, kwargs=kwargs
            )
            value, action = next(gen)
            result = None
            while True:
                if action == "result":
                    result = value
                try:
                    value, action = gen.send(value)
                except StopIteration:
                    break
            return cast("T", result)

    return wrapper  # type: ignore[return-value]


def dependency(*, scope_class: type[Scope] = NullScope) -> Callable[[TC], TC]:
    """
    Decorator to declare a dependency. You don't need to use it with default arguments,
    use it only if you want to change the scope of the dependency.
    Should be placed last in the decorator chain (on top).
    """

    if scope_class not in _scopes:
        _scopes[scope_class] = scope_class()

    def decorator(fn: TC) -> TC:
        _internal_registry.add(fn, scope_class=scope_class, override_scope=True)
        return fn

    return decorator


def init_dependencies() -> Awaitable:
    """
    Call this function to close dependencies. Usually, it should be called
    when your application is shutting down.
    """
    async_deps = []
    global_providers = _internal_registry.filter(
        lambda p: issubclass(p.scope_class, GlobalScope)
    )
    for provider in global_providers:
        resolver = LazyResolver(provider)
        value = resolver(provider.is_async)
        if provider.is_async:
            async_deps.append(value)

    if async_deps:
        return asyncio.gather(*async_deps)
    return NullAwaitable()


def shutdown_dependencies() -> Awaitable:
    """
    Call this function to close dependencies. Usually, it should be called
    when your application is shut down.
    """
    tasks = [scope.close_global() for scope in _scopes.values()]
    return asyncio.gather(*tasks)


class Dependency(NamedTuple):
    call: DependencyCallable

    def __call__(self) -> Dependency:
        return self

    def get_provider(self) -> Provider:
        return _internal_registry.get(self.call)


@dataclass(frozen=True)
class Provider:
    dependency: DependencyCallable
    is_async: bool
    scope_class: type[Scope]

    @classmethod
    def from_dependency(
        cls, dependency: DependencyCallable, scope_class: type[Scope]
    ) -> Provider:
        is_async = inspect.iscoroutinefunction(
            dependency
        ) or inspect.isasyncgenfunction(dependency)
        return cls(
            dependency=dependency,
            is_async=is_async,
            scope_class=scope_class,
        )

    def replace(self, scope_class: type[Scope] | None = None) -> Provider:
        kwargs = asdict(self)
        if scope_class is not None:
            kwargs["scope_class"] = scope_class
        return Provider(**kwargs)

    def get_scope(self) -> Scope:
        return _scopes[self.scope_class]

    def resolve_value(self, **kwargs: Any) -> Any:
        scope = self.get_scope()
        value_or_gen = self.dependency(**kwargs)
        if self.is_async:

            async def resolve_value_inner() -> Any:
                value_or_gen_ = value_or_gen
                if inspect.iscoroutine(value_or_gen):
                    value_or_gen_ = await value_or_gen_
                if inspect.isasyncgen(value_or_gen_):
                    context_manager = asynccontextmanager(
                        lambda *args, **kwargs: value_or_gen_
                    )
                    return await scope.exit_stack.enter_context(context_manager())
                return value_or_gen_

            return resolve_value_inner()

        if inspect.isgenerator(value_or_gen):
            context_manager = contextmanager(lambda *args, **kwargs: value_or_gen)
            return scope.exit_stack.enter_context(context_manager())
        return value_or_gen


def _wrapper_helper(
    dependant: DependNode,
    signature: inspect.Signature,
    is_async: bool,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Generator[Any, None, None]:
    bound = signature.bind(*args, **kwargs)
    bound.apply_defaults()
    arguments: dict[str, Any] = bound.arguments
    scopes: list[Scope] = []
    is_root = any(isinstance(value, Dependency) for value in bound.arguments.values())

    if is_root:
        arguments, scopes = _resolve_dependencies(dependant)

    for scope in scopes:
        scope.enter_decorator()
    for name, call in arguments.items():
        if isinstance(call, LazyResolver):
            value = yield call(is_async=is_async), "dependency"
            bound.arguments[name] = value

    try:
        result = dependant.value.call(*bound.args, **bound.kwargs)
    except Exception as e:
        for scope in scopes:
            scope.exit_decorator(e)
            yield scope.close_local(e), "close_scope"
        raise

    if inspect.isgenerator(result):

        @functools.wraps(result)  # type: ignore[arg-type]
        def gen() -> Generator[Any, None, None]:
            exception = None
            try:
                yield from result
            except Exception as e:
                exception = e
                raise
            finally:
                for scope in scopes:
                    scope.exit_decorator(exception)
                    scope.close_local(exception)

        yield gen(), "result"
        return
    elif inspect.isasyncgen(result):

        @functools.wraps(result)  # type: ignore[arg-type]
        async def gen() -> AsyncGenerator[Any, None]:
            exception = None
            try:
                async for item in result:
                    yield item
            except Exception as e:
                exception = e
                raise
            # TODO use
            #   https://flake8-async.readthedocs.io/en/latest/glossary.html#cancel-scope
            #   https://docs.python.org/3/library/asyncio-task.html#task-cancellation
            finally:
                for scope in scopes:
                    scope.exit_decorator(exception)
                    await scope.close_local(exception)  # noqa: ASYNC102

        yield gen(), "result"
        return

    yield result, "result"
    for scope in scopes:
        scope.exit_decorator()
        yield scope.close_local(), "close_scope"


def _resolve_dependencies(
    dependant: DependNode,
) -> tuple[dict[str, LazyResolver], list[Scope]]:
    scopes = set()
    resolved_dependencies = {}
    for dep in dependant.dependencies:
        values, dep_scopes = _resolve_dependencies(dep)
        resolved_dependencies.update(values)
        scopes.update(dep_scopes)

    if dependant.name is None:
        return resolved_dependencies, list(scopes)

    provider = dependant.value.get_provider()
    value = LazyResolver(
        provider=provider,
        kwargs=resolved_dependencies,
    )
    return {dependant.name: value}, [provider.get_scope()]


@dataclass
class DependNode:
    value: Dependency
    name: str | None
    dependencies: list[DependNode]


def _build_depend_tree(dependency: Dependency, name: str | None = None) -> DependNode:
    signature = inspect.signature(dependency.call)
    dependencies = []
    for name_, value in signature.parameters.items():
        param_dep = _extract_and_register_dependency_from_parameter(value)
        if param_dep is not None:
            dependencies.append(_build_depend_tree(param_dep, name=name_))
    return DependNode(value=dependency, dependencies=dependencies, name=name)


def _extract_and_register_dependency_from_parameter(
    value: inspect.Parameter,
) -> Dependency | None:
    if isinstance(value.default, Dependency):
        _internal_registry.add(value.default.call)
        return value.default

    if fastapi is None:
        return None  # type: ignore[unreachable]
    fastapi_dependency = None
    if isinstance(value.default, fastapi.params.Depends):
        fastapi_dependency = value.default.dependency
    elif get_origin(value.annotation) is Annotated:
        for metadata in value.annotation.__metadata__:
            if isinstance(metadata, fastapi.params.Depends):
                fastapi_dependency = metadata.dependency
                break
    if isinstance(fastapi_dependency, Dependency):
        _internal_registry.add(fastapi_dependency.call)  # type: ignore[unreachable]
        return fastapi_dependency
    return None


class LazyResolver:
    def __init__(
        self,
        provider: Provider,
        kwargs: dict[str, Any] | None = None,
    ) -> None:
        self.provider = provider
        self.kwargs = kwargs or {}

    def __call__(self, is_async: bool) -> Any:
        call = self._resolve_async if is_async else self._resolve
        return call()

    def _resolve(self) -> Any:
        scope = self.provider.get_scope()
        try:
            value = scope.get(self.provider.dependency)
        except KeyError:
            if self.provider.is_async:
                value = self.provider.dependency()
            else:
                with _lock:
                    try:
                        value = scope.get(self.provider.dependency)
                    except KeyError:
                        value = self.provider.resolve_value(**self.kwargs)
                        scope.set(self.provider.dependency, value)
        return value

    async def _resolve_async(self) -> Any:
        scope = self.provider.get_scope()
        try:
            value = scope.get(self.provider.dependency)
        except KeyError:
            with _lock:
                try:
                    value = scope.get(self.provider.dependency)
                except KeyError:
                    value = self.provider.resolve_value(**self.kwargs)
                    if self.provider.is_async:
                        value = await value
                    scope.set(self.provider.dependency, value)
        return value
