import inspect
import typing
from collections.abc import AsyncGenerator, Awaitable, Callable, Generator
from contextlib import AbstractAsyncContextManager, AbstractContextManager
from dataclasses import dataclass
from enum import Flag
from types import TracebackType
from typing import (
    ClassVar,
    Final,
    Self,
    final,
    override,
)

from wirio._service_lookup._async_factory_call_site import (
    AsyncFactoryCallSite,
)
from wirio._service_lookup._async_generator_factory_call_site import (
    AsyncGeneratorFactoryCallSite,
)
from wirio._service_lookup._call_site_visitor import CallSiteVisitor
from wirio._service_lookup._constant_call_site import (
    ConstantCallSite,
)
from wirio._service_lookup._constructor_call_site import (
    ConstructorCallSite,
)
from wirio._service_lookup._parameter_information import (
    ParameterInformation,
)
from wirio._service_lookup._sequence_call_site import SequenceCallSite
from wirio._service_lookup._service_call_site import (
    ServiceCallSite,
)
from wirio._service_lookup._service_provider_call_site import (
    ServiceProviderCallSite,
)
from wirio._service_lookup._supports_async_context_manager import (
    SupportsAsyncContextManager,
)
from wirio._service_lookup._supports_sync_context_manager import (
    SupportsSyncContextManager,
)
from wirio._service_lookup._sync_factory_call_site import (
    SyncFactoryCallSite,
)
from wirio._service_lookup._sync_generator_factory_call_site import (
    SyncGeneratorFactoryCallSite,
)
from wirio._service_lookup._typed_type import TypedType
from wirio.annotations import FromKeyedServicesInjectable
from wirio.exceptions import (
    CannotResolveParameterServiceFromImplementationFactoryError,
    GeneratorFactoryYieldedSeveralTimesError,
)
from wirio.service_provider_engine_scope import (
    ServiceProviderEngineScope,
)
from wirio.wirio_undefined import WirioUndefined


class _RuntimeResolverLock(Flag):
    NONE = 0
    SCOPE = 1
    ROOT = 2


@final
@dataclass(frozen=True)
class RuntimeResolverContext:
    scope: ServiceProviderEngineScope
    acquired_locks: _RuntimeResolverLock


@final
class _GeneratorFactoryDisposable(
    AbstractContextManager["_GeneratorFactoryDisposable"]
):
    _generator: Final[Generator[object]]
    _service: object | None

    def __init__(self, generator: Generator[object]) -> None:
        self._generator = generator
        self._service = None

    @property
    def service(self) -> object | None:
        return self._service

    @override
    def __enter__(self) -> Self:
        self._service = next(self._generator)
        return self

    @override
    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> bool | None:
        try:
            next(self._generator)
        except StopIteration:
            return None

        raise GeneratorFactoryYieldedSeveralTimesError


@final
class _AsyncGeneratorFactoryDisposable(
    AbstractAsyncContextManager["_AsyncGeneratorFactoryDisposable"]
):
    _generator: Final[AsyncGenerator[object]]
    _service: object | None

    def __init__(self, generator: AsyncGenerator[object]) -> None:
        self._generator = generator
        self._service = None

    @property
    def service(self) -> object | None:
        return self._service

    @override
    async def __aenter__(self) -> Self:
        self._service = await anext(self._generator)
        return self

    @override
    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> bool | None:
        try:
            await anext(self._generator)
        except StopAsyncIteration:
            return None

        raise GeneratorFactoryYieldedSeveralTimesError


@final
class CallSiteRuntimeResolver(CallSiteVisitor[RuntimeResolverContext, object | None]):
    INSTANCE: ClassVar["CallSiteRuntimeResolver"]

    async def resolve(
        self, call_site: ServiceCallSite, scope: ServiceProviderEngineScope
    ) -> object | None:
        # Fast path to avoid virtual calls if we already have the cached value in the root scope
        if scope.is_root_scope and call_site.value is not None:
            return call_site.value

        return await self._visit_call_site(
            call_site,
            RuntimeResolverContext(
                scope=scope, acquired_locks=_RuntimeResolverLock.NONE
            ),
        )

    @override
    async def _visit_call_site(
        self, call_site: ServiceCallSite, argument: RuntimeResolverContext
    ) -> object | None:
        overridden_call_site = self._get_overridden_call_site(call_site, argument)

        if not isinstance(overridden_call_site, WirioUndefined):
            return overridden_call_site

        return await super()._visit_call_site(call_site, argument)

    def _get_overridden_call_site(
        self, call_site: ServiceCallSite, argument: RuntimeResolverContext
    ) -> object | None | WirioUndefined:
        service_identifier = call_site.cache.key.service_identifier
        override_call_site = argument.scope.root_provider.get_overridden_call_site(
            service_identifier
        )

        if override_call_site is None:
            return WirioUndefined.INSTANCE

        return override_call_site.value

    @override
    async def _visit_root_cache(
        self, call_site: ServiceCallSite, argument: RuntimeResolverContext
    ) -> object | None:
        # If the value is already calculated, return it directly
        if call_site.value is not None:
            return call_site.value

        lock_type = _RuntimeResolverLock.ROOT
        service_provider_engine_scope = argument.scope.root_provider.root

        async with call_site.lock:
            # Lock the call_site and check if another coroutine already cached the value
            if call_site.value is not None:
                return call_site.value

            resolved_service = await self._visit_call_site_main(
                call_site=call_site,
                argument=RuntimeResolverContext(
                    scope=service_provider_engine_scope,
                    acquired_locks=argument.acquired_locks | lock_type,
                ),
            )
            await service_provider_engine_scope.capture_disposable(resolved_service)
            call_site.value = resolved_service
            return resolved_service

    @override
    async def _visit_scope_cache(
        self, call_site: ServiceCallSite, argument: RuntimeResolverContext
    ) -> object | None:
        # Check if we are in the situation where scoped service was promoted to singleton
        # and we need to lock the root
        if argument.scope.is_root_scope:
            return await self._visit_root_cache(call_site, argument)

        return await self._visit_cache(
            call_site, argument, argument.scope, _RuntimeResolverLock.SCOPE
        )

    async def _visit_cache(
        self,
        call_site: ServiceCallSite,
        argument: RuntimeResolverContext,
        service_provider_engine_scope: ServiceProviderEngineScope,
        lock_type: _RuntimeResolverLock,
    ) -> object | None:
        is_lock_taken = False
        resolved_services_lock = service_provider_engine_scope.resolved_services_lock
        resolved_services = service_provider_engine_scope.resolved_services

        # Taking locks only once allows us to fork resolution process
        # on another coroutine without causing the deadlock because we
        # always know that we are going to wait the other coroutine to finish before
        # releasing the lock
        if (argument.acquired_locks & lock_type) == _RuntimeResolverLock.NONE:
            await resolved_services_lock.acquire()
            is_lock_taken = True

        try:
            # Note: This method has already taken lock by the caller for resolution and access synchronization.
            # For scoped: takes a dictionary as both a resolution lock and a dictionary access lock.
            resolved_service = resolved_services.get(
                call_site.cache.key, WirioUndefined.INSTANCE
            )

            if resolved_service is not WirioUndefined.INSTANCE:
                return resolved_service

            resolved_service = await self._visit_call_site_main(
                call_site=call_site,
                argument=RuntimeResolverContext(
                    scope=service_provider_engine_scope,
                    acquired_locks=argument.acquired_locks | lock_type,
                ),
            )
            await service_provider_engine_scope.capture_disposable(resolved_service)
            resolved_services[call_site.cache.key] = resolved_service
            return resolved_service
        finally:
            if is_lock_taken:
                resolved_services_lock.release()

    @override
    async def _visit_dispose_cache(
        self, call_site: ServiceCallSite, argument: RuntimeResolverContext
    ) -> object | None:
        service = await self._visit_call_site_main(call_site, argument)
        return await argument.scope.capture_disposable(service)

    @override
    async def _visit_constructor(
        self,
        constructor_call_site: ConstructorCallSite,
        argument: RuntimeResolverContext,
    ) -> object:
        parameter_values: list[object | None] = []

        for parameter, parameter_call_site in zip(
            constructor_call_site.parameters,
            constructor_call_site.parameter_call_sites,
            strict=True,
        ):
            if parameter_call_site is None:
                if parameter.has_default_value:
                    parameter_values.append(parameter.default_value)
                    continue

                if parameter.is_optional:
                    parameter_values.append(None)
                    continue

                raise self._build_constructor_parameter_resolution_error(
                    parameter, constructor_call_site.service_type
                )

            parameter_service = await self._visit_call_site(
                parameter_call_site, argument
            )

            if parameter_service is None and not parameter.is_optional:
                raise self._build_constructor_parameter_resolution_error(
                    parameter, constructor_call_site.service_type
                )

            parameter_values.append(parameter_service)

        service = constructor_call_site.constructor_information.invoke(parameter_values)

        if isinstance(service, SupportsAsyncContextManager):
            await service.__aenter__()
        elif isinstance(service, SupportsSyncContextManager):
            service.__enter__()

        return service

    @override
    def _visit_constant(
        self,
        constant_call_site: ConstantCallSite,
        argument: RuntimeResolverContext,
    ) -> object:
        return constant_call_site.default_value

    @override
    async def _visit_sync_factory(
        self,
        sync_factory_call_site: SyncFactoryCallSite,
        argument: RuntimeResolverContext,
    ) -> object | None:
        parameter_services = await self.get_parameter_services(
            sync_factory_call_site.implementation_factory, argument.scope
        )
        service = sync_factory_call_site.implementation_factory(*parameter_services)

        if isinstance(service, SupportsAsyncContextManager):
            await service.__aenter__()
        elif isinstance(service, SupportsSyncContextManager):
            service.__enter__()

        return service

    @override
    async def _visit_async_factory(
        self,
        async_factory_call_site: AsyncFactoryCallSite,
        argument: RuntimeResolverContext,
    ) -> object | None:
        parameter_services = await self.get_parameter_services(
            async_factory_call_site.implementation_factory, argument.scope
        )
        service = await async_factory_call_site.implementation_factory(
            *parameter_services
        )

        if isinstance(service, SupportsAsyncContextManager):
            await service.__aenter__()
        elif isinstance(service, SupportsSyncContextManager):
            service.__enter__()

        return service

    @override
    async def _visit_sync_generator_factory(
        self,
        sync_generator_factory_call_site: SyncGeneratorFactoryCallSite,
        argument: RuntimeResolverContext,
    ) -> object | None:
        parameter_services = await self.get_parameter_services(
            sync_generator_factory_call_site.implementation_factory, argument.scope
        )
        service_generator = sync_generator_factory_call_site.implementation_factory(
            *parameter_services
        )
        disposable = _GeneratorFactoryDisposable(service_generator)
        disposable.__enter__()
        service = disposable.service
        await argument.scope.capture_disposable(disposable)
        return service

    @override
    async def _visit_async_generator_factory(
        self,
        async_generator_factory_call_site: AsyncGeneratorFactoryCallSite,
        argument: RuntimeResolverContext,
    ) -> object | None:
        parameter_services = await self.get_parameter_services(
            async_generator_factory_call_site.implementation_factory, argument.scope
        )
        service_generator = async_generator_factory_call_site.implementation_factory(
            *parameter_services
        )
        disposable = _AsyncGeneratorFactoryDisposable(service_generator)
        await disposable.__aenter__()
        service = disposable.service
        await argument.scope.capture_disposable(disposable)
        return service

    @override
    async def _visit_sequence(
        self, sequence_call_site: SequenceCallSite, argument: RuntimeResolverContext
    ) -> object | None:
        sequence: list[object | None] = []

        for service_call_site in sequence_call_site.service_call_sites:
            value = await self._visit_call_site(service_call_site, argument)
            sequence.append(value)

        return tuple(sequence)

    @override
    def _visit_service_provider(
        self,
        service_provider_call_site: ServiceProviderCallSite,
        argument: RuntimeResolverContext,
    ) -> object | None:
        return argument.scope

    async def get_parameter_services(
        self,
        implementation_factory: Callable[..., Awaitable[object]]
        | Callable[..., object],
        scope: ServiceProviderEngineScope,
    ) -> list[object | None]:
        parameter_services: list[object | None] = []
        signature = inspect.signature(implementation_factory)

        for parameter in signature.parameters.values():
            parameter_information = ParameterInformation(parameter)
            parameter_service: object | None = None

            if isinstance(
                parameter_information.injectable_dependency, FromKeyedServicesInjectable
            ):
                parameter_service = await scope.get_keyed_service_object(
                    parameter_information.injectable_dependency.key,
                    parameter_information.parameter_type,
                )
            else:
                parameter_service = await scope.get_service_object(
                    parameter_information.parameter_type
                )

            if parameter_service is None:
                if parameter_information.has_default_value:
                    parameter_services.append(parameter_information.default_value)
                elif parameter_information.is_optional:
                    parameter_services.append(None)
                else:
                    raise CannotResolveParameterServiceFromImplementationFactoryError(
                        parameter_information.parameter_type
                    )
            else:
                parameter_services.append(parameter_service)

        return parameter_services

    def _get_parameter_types(
        self,
        implementation_factory: Callable[..., Awaitable[object]]
        | Callable[..., object],
    ) -> list[type]:
        type_hints: dict[str, type] = typing.get_type_hints(implementation_factory)
        return [
            parameter_type
            for parameter_name, parameter_type in type_hints.items()
            if parameter_name != "return"
        ]

    def _build_constructor_parameter_resolution_error(
        self, parameter_information: ParameterInformation, service_type: TypedType
    ) -> RuntimeError:
        return RuntimeError(
            f"Unable to resolve service with type '{parameter_information.parameter_type}' while attempting to activate '{service_type}'"
        )


CallSiteRuntimeResolver.INSTANCE = CallSiteRuntimeResolver()
