import asyncio
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from typing import ClassVar, Final, final, override

from wirio._service_lookup._async_concurrent_dictionary import (
    AsyncConcurrentDictionary,
)
from wirio._service_lookup._async_factory_call_site import (
    AsyncFactoryCallSite,
)
from wirio._service_lookup._call_site_chain import CallSiteChain
from wirio._service_lookup._constant_call_site import (
    ConstantCallSite,
)
from wirio._service_lookup._constructor_call_site import (
    ConstructorCallSite,
)
from wirio._service_lookup._constructor_information import (
    ConstructorInformation,
)
from wirio._service_lookup._parameter_information import (
    ParameterInformation,
)
from wirio._service_lookup._result_cache import ResultCache
from wirio._service_lookup._service_call_site import (
    ServiceCallSite,
)
from wirio._service_lookup._service_identifier import (
    ServiceIdentifier,
)
from wirio._service_lookup._sync_factory_call_site import (
    SyncFactoryCallSite,
)
from wirio._service_lookup._typed_type import TypedType
from wirio._service_lookup.service_cache_key import ServiceCacheKey
from wirio.abstractions.base_service_provider import BaseServiceProvider
from wirio.abstractions.keyed_service import KeyedService
from wirio.abstractions.service_key_lookup_mode import (
    ServiceKeyLookupMode,
)
from wirio.abstractions.service_provider_is_keyed_service import (
    ServiceProviderIsKeyedService,
)
from wirio.abstractions.service_provider_is_service import (
    ServiceProviderIsService,
)
from wirio.abstractions.service_scope_factory import (
    ServiceScopeFactory,
)
from wirio.annotations import (
    FromKeyedServicesInjectable,
    ServiceKeyInjectable,
)
from wirio.exceptions import (
    CannotResolveServiceError,
    InvalidServiceDescriptorError,
    InvalidServiceKeyTypeError,
)
from wirio.service_descriptor import ServiceDescriptor


@final
class _ServiceDescriptorCacheItem:
    _item: ServiceDescriptor | None
    _items: list[ServiceDescriptor] | None

    def __init__(self) -> None:
        self._item = None
        self._items = None

    @property
    def last(self) -> ServiceDescriptor:
        if self._items is not None and len(self._items) > 0:
            return self._items[len(self._items) - 1]

        assert self._item is not None
        return self._item

    def add(self, descriptor: ServiceDescriptor) -> "_ServiceDescriptorCacheItem":
        new_cache_item = _ServiceDescriptorCacheItem()

        if self._item is None:
            new_cache_item._item = descriptor
        else:
            new_cache_item._item = self._item
            new_cache_item._items = self._items if self._items is not None else []
            new_cache_item._items.append(descriptor)

        return new_cache_item


@dataclass(frozen=True)
class _ServiceOverride:
    exists: bool
    value: object | None = None


@final
class CallSiteFactory(ServiceProviderIsKeyedService, ServiceProviderIsService):
    _DEFAULT_SLOT: ClassVar[int] = 0

    _descriptors: Final[list[ServiceDescriptor]]
    _descriptor_lookup: Final[dict[ServiceIdentifier, _ServiceDescriptorCacheItem]]
    _call_site_cache: Final[AsyncConcurrentDictionary[ServiceCacheKey, ServiceCallSite]]
    _call_site_locks: Final[AsyncConcurrentDictionary[ServiceIdentifier, asyncio.Lock]]
    _service_overrides: Final[dict[ServiceIdentifier, list[object | None]]]

    def __init__(self, descriptors: list["ServiceDescriptor"]) -> None:
        self._descriptors = descriptors.copy()
        self._descriptor_lookup = {}
        self._call_site_cache = AsyncConcurrentDictionary[
            ServiceCacheKey, ServiceCallSite
        ]()
        self._call_site_locks = AsyncConcurrentDictionary[
            ServiceIdentifier, asyncio.Lock
        ]()
        self._service_overrides = {}
        self._populate()

    @override
    def is_service(self, service_type: type) -> bool:
        return self._is_service(
            ServiceIdentifier.from_service_type(
                service_type=TypedType.from_type(service_type)
            )
        )

    @override
    def is_keyed_service(self, service_key: object | None, service_type: type) -> bool:
        return self._is_service(
            ServiceIdentifier.from_service_type(
                service_type=TypedType.from_type(service_type), service_key=service_key
            )
        )

    async def get_call_site(
        self, service_identifier: ServiceIdentifier, call_site_chain: CallSiteChain
    ) -> ServiceCallSite | None:
        overridden_call_site = self._get_overridden_call_site(service_identifier)

        if overridden_call_site is not None:
            return overridden_call_site

        service_cache_key = ServiceCacheKey(service_identifier, self._DEFAULT_SLOT)
        service_call_site = self._call_site_cache.get(service_cache_key)

        if service_call_site is None:
            return await self._create_call_site(
                service_identifier=service_identifier, call_site_chain=call_site_chain
            )

        return service_call_site

    async def add(
        self, service_identifier: ServiceIdentifier, service_call_site: ServiceCallSite
    ) -> None:
        cache_key = ServiceCacheKey(service_identifier, self._DEFAULT_SLOT)
        await self._call_site_cache.upsert(key=cache_key, value=service_call_site)

    @contextmanager
    def override_service(
        self,
        service_identifier: ServiceIdentifier,
        implementation_instance: object | None,
    ) -> Generator[None]:
        self._add_override(service_identifier, implementation_instance)

        try:
            yield
        finally:
            self._remove_override(service_identifier)

    def get_overridden_call_site(
        self, service_identifier: ServiceIdentifier
    ) -> ServiceCallSite | None:
        return self._get_overridden_call_site(service_identifier)

    async def _create_call_site(
        self, service_identifier: ServiceIdentifier, call_site_chain: CallSiteChain
    ) -> ServiceCallSite | None:
        async def _create_new_lock(_: ServiceIdentifier) -> asyncio.Lock:
            return asyncio.Lock()

        # We need to lock the resolution process for a single service type at a time.
        # Consider the following:
        # C -> D -> A
        # E -> D -> A
        # Resolving C and E in parallel means that they will be modifying the callsite cache concurrently
        # to add the entry for C and E, but the resolution of D and A is synchronized
        # to make sure C and E both reference the same instance of the callsite.
        #
        # This is to make sure we can safely store singleton values on the callsites themselves

        call_site_lock = await self._call_site_locks.get_or_add(
            service_identifier, _create_new_lock
        )

        # Check if the lock is already acquired to prevent deadlocks in case of re-entrancy
        if call_site_lock.locked():
            call_site_chain.check_circular_dependency(service_identifier)

        async with call_site_lock:
            return await self._try_create_exact_from_service_identifier(
                service_identifier, call_site_chain
            )

    def _populate(self) -> None:
        for descriptor in self._descriptors:
            cache_key = ServiceIdentifier.from_descriptor(descriptor)
            cache_item = self._descriptor_lookup.get(
                cache_key, _ServiceDescriptorCacheItem()
            )
            self._descriptor_lookup[cache_key] = cache_item.add(descriptor)

    async def _try_create_exact_from_service_identifier(
        self, service_identifier: ServiceIdentifier, call_site_chain: CallSiteChain
    ) -> ServiceCallSite | None:
        service_descriptor_cache_item = self._descriptor_lookup.get(
            service_identifier, None
        )

        if service_descriptor_cache_item is not None:
            return await self._try_create_exact_from_service_descriptor(
                service_descriptor_cache_item.last,
                service_identifier,
                call_site_chain,
                self._DEFAULT_SLOT,
            )

        # Check if there is a registration with `KeyedService.ANY_KEY`
        if service_identifier.service_key is not None:
            catch_all_identifier = ServiceIdentifier(
                service_type=service_identifier.service_type,
                service_key=KeyedService.ANY_KEY,
            )

            service_descriptor_cache_item = self._descriptor_lookup.get(
                catch_all_identifier, None
            )

            if service_descriptor_cache_item is not None:
                return await self._try_create_exact_from_service_descriptor(
                    service_descriptor_cache_item.last,
                    service_identifier,
                    call_site_chain,
                    self._DEFAULT_SLOT,
                )

        return None

    async def _try_create_exact_from_service_descriptor(
        self,
        service_descriptor: ServiceDescriptor,
        service_identifier: ServiceIdentifier,
        call_site_chain: CallSiteChain,
        slot: int,
    ) -> ServiceCallSite | None:
        if not self._should_create_exact(
            service_descriptor.service_type, service_identifier.service_type
        ):
            return None

        return await self._create_exact(
            service_descriptor, service_identifier, call_site_chain, slot
        )

    def _should_create_exact(
        self, descriptor_type: TypedType, service_type: TypedType
    ) -> bool:
        return descriptor_type == service_type

    async def _create_exact(
        self,
        service_descriptor: ServiceDescriptor,
        service_identifier: ServiceIdentifier,
        call_site_chain: CallSiteChain,
        slot: int,
    ) -> ServiceCallSite:
        call_site_key = ServiceCacheKey(service_identifier, slot)
        service_call_site = self._call_site_cache.get(call_site_key)

        if service_call_site is not None:
            return service_call_site

        cache = ResultCache.from_lifetime(
            service_descriptor.lifetime, service_identifier, slot
        )

        if service_descriptor.has_implementation_instance():
            service_call_site = ConstantCallSite(
                service_type=service_descriptor.service_type,
                default_value=service_descriptor.get_implementation_instance(),
                service_key=service_descriptor.service_key,
            )
        elif (
            not service_descriptor.is_keyed_service
            and service_descriptor.sync_implementation_factory is not None
        ):
            service_call_site = SyncFactoryCallSite.from_implementation_factory(
                cache=cache,
                service_type=service_descriptor.service_type,
                implementation_factory=service_descriptor.sync_implementation_factory,
            )
        elif (
            service_descriptor.is_keyed_service
            and service_descriptor.keyed_sync_implementation_factory is not None
        ):
            service_call_site = SyncFactoryCallSite.from_keyed_implementation_factory(
                cache=cache,
                service_type=service_descriptor.service_type,
                implementation_factory=service_descriptor.keyed_sync_implementation_factory,
                service_key=service_identifier.service_key,
            )
        elif (
            not service_descriptor.is_keyed_service
            and service_descriptor.async_implementation_factory is not None
        ):
            service_call_site = AsyncFactoryCallSite.from_implementation_factory(
                cache=cache,
                service_type=service_descriptor.service_type,
                implementation_factory=service_descriptor.async_implementation_factory,
            )
        elif (
            service_descriptor.is_keyed_service
            and service_descriptor.keyed_async_implementation_factory is not None
        ):
            service_call_site = AsyncFactoryCallSite.from_keyed_implementation_factory(
                cache=cache,
                service_type=service_descriptor.service_type,
                implementation_factory=service_descriptor.keyed_async_implementation_factory,
                service_key=service_identifier.service_key,
            )
        elif service_descriptor.has_implementation_type():
            implementation_type = service_descriptor.get_implementation_type()
            assert implementation_type is not None
            service_call_site = await self._create_constructor_call_site(
                cache=cache,
                service_identifier=service_identifier,
                implementation_type=implementation_type,
                call_site_chain=call_site_chain,
            )
        else:
            raise InvalidServiceDescriptorError

        await self._call_site_cache.upsert(key=call_site_key, value=service_call_site)
        return service_call_site

    def _get_overridden_call_site(
        self, service_identifier: ServiceIdentifier
    ) -> ServiceCallSite | None:
        override = self._get_overridden_instance(service_identifier)

        if not override.exists:
            return None

        return ConstantCallSite(
            service_type=service_identifier.service_type,
            default_value=override.value,
            service_key=service_identifier.service_key,
        )

    def _get_overridden_instance(
        self, service_identifier: ServiceIdentifier
    ) -> _ServiceOverride:
        overrides = self._service_overrides.get(service_identifier)

        if overrides is not None and len(overrides) > 0:
            return _ServiceOverride(exists=True, value=overrides[-1])

        catch_all_identifier = self._get_catch_all_service_identifier(
            service_identifier
        )

        if catch_all_identifier is None:
            return _ServiceOverride(exists=False, value=None)

        overrides = self._service_overrides.get(catch_all_identifier)

        if overrides is not None and len(overrides) > 0:
            return _ServiceOverride(exists=True, value=overrides[-1])

        return _ServiceOverride(exists=False, value=None)

    def _get_catch_all_service_identifier(
        self, service_identifier: ServiceIdentifier
    ) -> ServiceIdentifier | None:
        if (
            service_identifier.service_key is None
            or service_identifier.service_key == KeyedService.ANY_KEY
        ):
            return None

        return ServiceIdentifier.from_service_type(
            service_type=service_identifier.service_type,
            service_key=KeyedService.ANY_KEY,
        )

    def _add_override(
        self,
        service_identifier: ServiceIdentifier,
        implementation_instance: object | None,
    ) -> None:
        overrides = self._service_overrides.setdefault(service_identifier, [])
        overrides.append(implementation_instance)

    def _remove_override(self, service_identifier: ServiceIdentifier) -> None:
        overrides = self._service_overrides.get(service_identifier)
        assert overrides is not None
        overrides.pop()

        if len(overrides) == 0:
            self._service_overrides.pop(service_identifier)

    async def _create_constructor_call_site(
        self,
        cache: ResultCache,
        service_identifier: ServiceIdentifier,
        implementation_type: TypedType,
        call_site_chain: CallSiteChain,
    ) -> ServiceCallSite:
        try:
            call_site_chain.add(service_identifier, implementation_type)
            parameter_call_sites: list[ServiceCallSite | None] | None = None
            constructor_information = ConstructorInformation(implementation_type)
            parameters = constructor_information.get_parameters()
            parameter_call_sites = await self._create_argument_call_sites(
                service_identifier=service_identifier,
                implementation_type=implementation_type,
                parameters=parameters,
                call_site_chain=call_site_chain,
            )
            return ConstructorCallSite(
                cache=cache,
                service_type=service_identifier.service_type,
                constructor_information=constructor_information,
                parameters=parameters,
                parameter_call_sites=parameter_call_sites,
                service_key=service_identifier.service_key,
            )
        finally:
            call_site_chain.remove(service_identifier)

    async def _create_argument_call_sites(  # noqa: C901, PLR0912
        self,
        service_identifier: ServiceIdentifier,
        implementation_type: TypedType,
        parameters: list[ParameterInformation],
        call_site_chain: CallSiteChain,
    ) -> list[ServiceCallSite | None]:
        if len(parameters) == 0:
            return []

        parameter_call_sites: list[ServiceCallSite | None] = []

        for parameter in parameters:
            call_site: ServiceCallSite | None = None
            is_keyed_parameter = False
            parameter_type = parameter.parameter_type

            if parameter.injectable_dependency is not None:
                if service_identifier.service_key is not None and isinstance(
                    parameter.injectable_dependency, ServiceKeyInjectable
                ):
                    # Even though the parameter may be strongly typed, support `object` if `ANY_KEY` is used

                    if service_identifier.service_key == KeyedService.ANY_KEY:
                        parameter_type = TypedType.from_type(object)
                    elif parameter_type.to_type() is not type(
                        service_identifier.service_key
                    ) and parameter_type.to_type() is not type(object):
                        raise InvalidServiceKeyTypeError

                    call_site = ConstantCallSite(
                        service_type=parameter_type,
                        default_value=service_identifier.service_key,
                    )
                elif isinstance(
                    parameter.injectable_dependency, FromKeyedServicesInjectable
                ):
                    service_key: object | None = None

                    match parameter.injectable_dependency.lookup_mode:
                        case ServiceKeyLookupMode.INHERIT_KEY:
                            service_key = service_identifier.service_key
                        case ServiceKeyLookupMode.EXPLICIT_KEY:
                            service_key = parameter.injectable_dependency.key
                        case ServiceKeyLookupMode.NULL_KEY:
                            service_key = None

                    if service_key is not None:
                        call_site = await self.get_call_site(
                            ServiceIdentifier.from_service_type(
                                service_type=parameter_type, service_key=service_key
                            ),
                            call_site_chain=call_site_chain,
                        )
                        is_keyed_parameter = True

            if not is_keyed_parameter and call_site is None:
                call_site = await self.get_call_site(
                    ServiceIdentifier.from_service_type(parameter_type), call_site_chain
                )

            if call_site is None and parameter.has_default_value:
                if parameter.is_optional:
                    parameter_call_sites.append(None)
                    continue

                call_site = ConstantCallSite(
                    service_type=parameter_type,
                    default_value=parameter.default_value,
                )

            if call_site is None and parameter.is_optional:
                parameter_call_sites.append(None)
                continue

            if call_site is None:
                raise CannotResolveServiceError(
                    parameter_type=parameter_type,
                    implementation_type=implementation_type,
                )

            parameter_call_sites.append(call_site)

        return parameter_call_sites

    def _is_service(self, service_identifier: ServiceIdentifier) -> bool:
        service_type = service_identifier.service_type

        if service_identifier in self._descriptor_lookup:
            return True

        if (
            service_identifier.service_key is not None
            and ServiceIdentifier(
                service_type=service_type, service_key=KeyedService.ANY_KEY
            )
            in self._descriptor_lookup
        ):
            return True

        return (
            service_type == TypedType.from_type(BaseServiceProvider)
            or service_type == TypedType.from_type(ServiceScopeFactory)
            or service_type == TypedType.from_type(ServiceProviderIsService)
            or service_type == TypedType.from_type(ServiceProviderIsKeyedService)
        )
