from collections.abc import Callable, Coroutine
from typing import TypeVar

from ..types import Message, Context
from .base import Injector, InjectorWithCleanup

__all__ = [
    "Lazy",
    "MessageInjector",
    "ContextInjector",
]


class MessageInjector(Injector[Message]):
    async def inject(self, message: Message, context: Context) -> Message:
        return message

MessageInjector = MessageInjector()


class ContextInjector(Injector[Context]):
    async def inject(self, message: Message, context: Context) -> Context:
        return context

ContextInjector = ContextInjector()

Dep = TypeVar("Dep")


class Lazy(InjectorWithCleanup[Callable[[], Coroutine[None, None, Dep]]]):
    def __init__(self, deferred: Injector[Dep]):
        self.deferred = deferred

    async def inject(
        self, message: Message, context: Context
    ) -> Callable[[], Coroutine[None, None, Dep]]:
        class _Wrapper:
            def __init__(self, deferred):
                self._calculated = None
                async def call():
                    if self._calculated is None:
                        self._calculated = await deferred.inject(message, context)
                    return self._calculated
                self._call = call

            def __call__(self):
                return self._call()
        
        return _Wrapper(self.deferred)

    async def cleanup(self, dep: Callable[[], Coroutine[None, None, Dep]]):
        if isinstance(self.deferred, InjectorWithCleanup) and dep._calculated is not None:
            await self.deferred.cleanup(dep._calculated)