"""Contextual abstractions for dependency injection using contextvars."""

from __future__ import annotations

import asyncio
import inspect
from collections import ChainMap
from collections.abc import AsyncIterator, Iterator
from contextlib import asynccontextmanager, contextmanager
from contextvars import ContextVar
from contextvars import Token as ContextToken
from types import TracebackType
from typing import Any, Awaitable, Callable, TypeVar, cast
from weakref import WeakValueDictionary

from .exceptions import AsyncCleanupRequiredError
from .tokens import Scope, Token

__all__ = [
    "ContextualContainer",
    "RequestScope",
    "SessionScope",
    "get_current_context",
    "set_context",
]

T = TypeVar("T")

# Global context variable for DI scopes
_context_stack: ContextVar[ChainMap[Token[object], object] | None] = ContextVar(
    "pyinj_context_stack", default=None
)

# Session context for longer-lived dependencies
_session_context: ContextVar[dict[Token[object], object] | None] = ContextVar(
    "pyinj_session_context", default=None
)

# Per-scope explicit cleanup stacks (registered via container.register_context)
_request_cleanup_sync: ContextVar[list[Callable[[], None]] | None] = ContextVar(
    "pyinj_request_cleanup_sync", default=None
)
_request_cleanup_async: ContextVar[
    list[Callable[[], Awaitable[None]]] | None
] = ContextVar("pyinj_request_cleanup_async", default=None)

_session_cleanup_sync: ContextVar[list[Callable[[], None]] | None] = ContextVar(
    "pyinj_session_cleanup_sync", default=None
)
_session_cleanup_async: ContextVar[
    list[Callable[[], Awaitable[None]]] | None
] = ContextVar("pyinj_session_cleanup_async", default=None)


def get_current_context() -> ChainMap[Token[object], object] | None:
    """Get current dependency context."""
    return _context_stack.get()


def set_context(
    context: ChainMap[Token[object], object],
) -> ContextToken[ChainMap[Token[object], object] | None]:
    """
    Set the current dependency context.

    Args:
        context: ChainMap of dependency caches

    Returns:
        Token for resetting context
    """
    return _context_stack.set(context)


class ContextualContainer:
    """Base container adding request/session context via ``contextvars``.

    Context flows implicitly across awaits; request/session lifetimes
    are enforced by the :class:`ScopeManager`.
    """

    def __init__(self) -> None:
        """Initialize contextual container."""
        # Singleton cache (process-wide)
        self._singletons: dict[Token[object], object] = {}

        # Weak cache for transients (auto-cleanup)
        self._transients: WeakValueDictionary[Token[object], object] = (
            WeakValueDictionary()
        )

        # Providers registry (value type depends on concrete container)
        self._providers: dict[Token[object], Any] = {}

        # Async locks for thread-safe singleton creation
        self._async_locks: dict[Token[object], asyncio.Lock] = {}

        # Track resources for cleanup
        from .protocols import SupportsAsyncClose, SupportsClose

        self._resources: list[SupportsClose | SupportsAsyncClose] = []

        # Scope manager (RAII contexts, precedence enforcement)
        self._scope_manager = ScopeManager(self)

    # ---- Registration of per-scope cleanup operations (used by Container) ----

    def _register_request_cleanup_sync(self, fn: Callable[[], None]) -> None:
        """Register a sync cleanup for the current request scope.

        Internal API called by the container when a sync context-managed resource
        is entered within a request scope. Cleanups run in LIFO order on scope exit.
        """
        stack = _request_cleanup_sync.get()
        if stack is None:
            raise RuntimeError("No active request scope for registering cleanup")
        stack.append(fn)

    def _register_request_cleanup_async(
        self, fn: Callable[[], Awaitable[None]]
    ) -> None:
        """Register an async cleanup for the current request scope.

        Internal API called by the container for async context-managed resources.
        Cleanups run before sync cleanups on async scope exit.
        """
        stack = _request_cleanup_async.get()
        if stack is None:
            raise RuntimeError("No active request scope for registering async cleanup")
        stack.append(fn)

    def _register_session_cleanup_sync(self, fn: Callable[[], None]) -> None:
        """Register a sync cleanup for the active session scope.

        Internal API used for session-scoped sync context-managed resources.
        """
        stack = _session_cleanup_sync.get()
        if stack is None:
            raise RuntimeError("No active session scope for registering cleanup")
        stack.append(fn)

    def _register_session_cleanup_async(
        self, fn: Callable[[], Awaitable[None]]
    ) -> None:
        """Register an async cleanup for the active session scope.

        Internal API used for session-scoped async context-managed resources.
        """
        stack = _session_cleanup_async.get()
        if stack is None:
            raise RuntimeError("No active session scope for registering async cleanup")
        stack.append(fn)

    def _put_in_current_request_cache(self, token: Token[T], instance: T) -> None:
        """Insert a value into the current request cache unconditionally.

        This bypasses scope checks and is intended for temporary overrides
        that should only affect the current context.
        """
        context = _context_stack.get()
        if context is not None and hasattr(context, "maps") and len(context.maps) > 0:
            # The top-most map holds request-local values
            context.maps[0][cast(Token[object], token)] = cast(object, instance)

    @contextmanager
    def request_scope(self) -> Iterator[ContextualContainer]:
        """Create a request scope (similar to a web request lifecycle).

        Example:
            with container.request_scope():
                service = container.get(ServiceToken)

        Yields:
            Self for chaining.
        """
        with self._scope_manager.request_scope():
            yield self

    @asynccontextmanager
    async def async_request_scope(self) -> AsyncIterator[ContextualContainer]:
        """Async context manager variant of :meth:`request_scope`.

        Example:
            async with container.async_request_scope():
                service = await container.aget(ServiceToken)
        """
        async with self._scope_manager.async_request_scope():
            yield self

    @contextmanager
    def session_scope(self) -> Iterator[ContextualContainer]:
        """
        Create a session scope (longer-lived than request).

        Session scopes persist across multiple requests but are
        isolated between different sessions (e.g., users).
        """
        with self._scope_manager.session_scope():
            yield self

    def _cleanup_scope(self, cache: dict[Token[object], Any]) -> None:
        """
        Clean up resources in LIFO order.

        Args:
            cache: Cache of resources to clean up
        """
        for resource in reversed(list(cache.values())):
            try:
                # Early fail if async cleanup is required in a sync context
                aclose = getattr(resource, "aclose", None)
                aexit = getattr(resource, "__aexit__", None)
                supports_sync = hasattr(resource, "close") or hasattr(
                    resource, "__exit__"
                )
                if (
                    (aclose and inspect.iscoroutinefunction(aclose))
                    or (aexit and inspect.iscoroutinefunction(aexit))
                ) and not supports_sync:
                    raise AsyncCleanupRequiredError(
                        type(resource).__name__,
                        "Use an async request/session scope.",
                    )
                close = getattr(resource, "close", None)
                if close is not None and inspect.iscoroutinefunction(close):
                    raise AsyncCleanupRequiredError(
                        type(resource).__name__,
                        "Use an async request/session scope.",
                    )
                # Prefer context-manager exit over ad-hoc close for sync path
                if hasattr(resource, "__exit__"):
                    resource.__exit__(None, None, None)
                elif close is not None:
                    close()
            except AsyncCleanupRequiredError:
                raise
            except Exception:
                # Log but don't fail cleanup
                pass

    async def _async_cleanup_scope(self, cache: dict[Token[object], Any]) -> None:
        """
        Async cleanup of resources.

        Args:
            cache: Cache of resources to clean up
        """
        tasks: list[Any] = []
        loop = asyncio.get_running_loop()

        for resource in reversed(list(cache.values())):
            aclose = getattr(resource, "aclose", None)
            if aclose and callable(aclose):
                res = aclose()
                if inspect.isawaitable(res):
                    tasks.append(res)
                    continue
            aexit = getattr(resource, "__aexit__", None)
            if aexit and callable(aexit):
                res = aexit(None, None, None)
                if inspect.isawaitable(res):
                    tasks.append(res)
                    continue
            close = getattr(resource, "close", None)
            if close:
                if inspect.iscoroutinefunction(close):
                    tasks.append(close())
                else:
                    tasks.append(loop.run_in_executor(None, close))

        if tasks:
            # Gather with return_exceptions to prevent one failure from stopping others
            await asyncio.gather(*tasks, return_exceptions=True)

    def resolve_from_context(self, token: Token[T]) -> T | None:
        """
        Resolve dependency from current context.

        Args:
            token: Token to resolve

        Returns:
            Resolved instance or None if not in context
        """
        return self._scope_manager.resolve_from_context(token)

    def store_in_context(self, token: Token[T], instance: T) -> None:
        """
        Store instance in appropriate context.

        Args:
            token: Token for the instance
            instance: Instance to store
        """
        self._scope_manager.store_in_context(token, instance)

    def clear_request_context(self) -> None:
        """Clear current request context."""
        self._scope_manager.clear_request_context()

    def clear_session_context(self) -> None:
        """Clear current session context."""
        self._scope_manager.clear_session_context()

    def clear_all_contexts(self) -> None:
        """Clear all contexts including singletons."""
        self._scope_manager.clear_all_contexts()


class ScopeManager:
    """Scope orchestration with RAII managers and explicit precedence.

    Precedence: REQUEST > SESSION > SINGLETON. Uses ContextVars for async safety.
    """

    def __init__(self, container: ContextualContainer) -> None:
        self._container = container

    @contextmanager
    def request_scope(self) -> Iterator[None]:
        request_cache: dict[Token[object], object] = {}
        current = _context_stack.get()
        if current is None:
            new_context = ChainMap(request_cache, self._container._singletons)
        else:
            new_context = ChainMap(request_cache, *current.maps)
        token = _context_stack.set(new_context)
        # Initialize per-request cleanup stacks
        req_sync_token = _request_cleanup_sync.set([])
        req_async_token = _request_cleanup_async.set([])
        try:
            yield
        finally:
            # Cleanup resources stored via context introspection (legacy path)
            self._container._cleanup_scope(request_cache)
            # Run explicit per-request cleanup stacks (LIFO)
            try:
                sync_fns = _request_cleanup_sync.get() or []
                for fn in reversed(sync_fns):
                    try:
                        fn()
                    except Exception:
                        pass
            finally:
                _request_cleanup_sync.reset(req_sync_token)
            _request_cleanup_async.reset(req_async_token)
            _context_stack.reset(token)

    @asynccontextmanager
    async def async_request_scope(self) -> AsyncIterator[None]:
        request_cache: dict[Token[object], object] = {}
        current = _context_stack.get()
        if current is None:
            new_context = ChainMap(request_cache, self._container._singletons)
        else:
            new_context = ChainMap(request_cache, *current.maps)
        token = _context_stack.set(new_context)
        # Initialize per-request cleanup stacks
        req_sync_token = _request_cleanup_sync.set([])
        req_async_token = _request_cleanup_async.set([])
        try:
            yield
        finally:
            # Async cleanup for resources stored via context introspection
            await self._container._async_cleanup_scope(request_cache)
            # Then run explicit per-request async cleanups (LIFO)
            async_fns = _request_cleanup_async.get() or []
            if async_fns:
                await asyncio.gather(
                    *[fn() for fn in reversed(async_fns)], return_exceptions=True
                )
            # Finally run any sync cleanups
            sync_fns = _request_cleanup_sync.get() or []
            for fn in reversed(sync_fns):
                try:
                    fn()
                except Exception:
                    pass
            _request_cleanup_sync.reset(req_sync_token)
            _request_cleanup_async.reset(req_async_token)
            _context_stack.reset(token)

    @contextmanager
    def session_scope(self) -> Iterator[None]:
        existing = _session_context.get()
        if existing is None:
            session_cache: dict[Token[object], object] = {}
            session_token = _session_context.set(session_cache)
            # Only create new cleanup stacks when opening a new session
            sess_sync_token = _session_cleanup_sync.set([])
            sess_async_token = _session_cleanup_async.set([])
        else:
            session_cache = existing
            session_token = None
            sess_sync_token = None
            sess_async_token = None
        current = _context_stack.get()
        if current is None:
            new_context = ChainMap(session_cache, self._container._singletons)
        else:
            new_context = ChainMap(
                current.maps[0], session_cache, self._container._singletons
            )
        context_token = _context_stack.set(new_context)
        try:
            yield
        finally:
            _context_stack.reset(context_token)
            if session_token:
                # Run explicit per-session cleanup stacks (LIFO) when session ends
                try:
                    sync_fns = _session_cleanup_sync.get() or []
                    for fn in reversed(sync_fns):
                        try:
                            fn()
                        except Exception:
                            pass
                finally:
                    if sess_sync_token is not None:
                        _session_cleanup_sync.reset(sess_sync_token)
                if sess_async_token is not None:
                    _session_cleanup_async.reset(sess_async_token)
                _session_context.reset(session_token)

    def resolve_from_context(self, token: Token[T]) -> T | None:
        context = _context_stack.get()
        if context is not None:
            key = cast(Token[object], token)
            if key in context:
                return cast(T, context[key])
        if token.scope == Scope.SESSION:
            session = _session_context.get()
            if session and token in session:
                return cast(T, session[cast(Token[object], token)])
        if token.scope == Scope.SINGLETON and token in self._container._singletons:
            return cast(T, self._container._singletons[cast(Token[object], token)])
        if token.scope == Scope.TRANSIENT and token in self._container._transients:
            return cast(T, self._container._transients[cast(Token[object], token)])
        return None

    def store_in_context(self, token: Token[T], instance: T) -> None:
        if token.scope == Scope.SINGLETON:
            self._container._singletons[cast(Token[object], token)] = cast(
                object, instance
            )
        elif token.scope == Scope.REQUEST:
            self._container._put_in_current_request_cache(token, instance)
        elif token.scope == Scope.SESSION:
            session = _session_context.get()
            if session is not None:
                session[cast(Token[object], token)] = cast(object, instance)
        elif token.scope == Scope.TRANSIENT:
            try:
                self._container._transients[cast(Token[object], token)] = cast(
                    object, instance
                )
            except TypeError:
                pass

    def clear_request_context(self) -> None:
        context = _context_stack.get()
        if context is not None and hasattr(context, "maps") and len(context.maps) > 0:
            context.maps[0].clear()

    def clear_session_context(self) -> None:
        session = _session_context.get()
        if session is not None:
            session.clear()

    def clear_all_contexts(self) -> None:
        self._container._singletons.clear()
        self._container._transients.clear()
        self.clear_request_context()
        self.clear_session_context()


class RequestScope:
    """
    Helper class for request-scoped dependencies.

    Example:
        async with RequestScope(container) as scope:
            service = scope.resolve(ServiceToken)
    """

    def __init__(self, container: ContextualContainer):
        """Initialize request scope."""
        self.container = container
        self._context_manager = None
        self._async_context_manager = None

    def __enter__(self) -> RequestScope:
        """Enter request scope."""
        self._context_manager = self.container.request_scope()
        self._context_manager.__enter__()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Exit request scope."""
        if self._context_manager:
            self._context_manager.__exit__(exc_type, exc_val, exc_tb)

    async def __aenter__(self) -> RequestScope:
        """Async enter request scope."""
        self._async_context_manager = self.container.async_request_scope()
        await self._async_context_manager.__aenter__()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Async exit request scope."""
        if self._async_context_manager:
            await self._async_context_manager.__aexit__(exc_type, exc_val, exc_tb)

    def resolve(self, token: Token[T]) -> T | None:
        """Resolve dependency in this scope."""
        return self.container.resolve_from_context(token)


class SessionScope:
    """
    Helper class for session-scoped dependencies.

    Example:
        with SessionScope(container) as scope:
            user = scope.resolve(UserToken)
    """

    def __init__(self, container: ContextualContainer):
        """Initialize session scope."""
        self.container = container
        self._context_manager = None

    def __enter__(self) -> SessionScope:
        """Enter session scope."""
        self._context_manager = self.container.session_scope()
        self._context_manager.__enter__()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Exit session scope."""
        if self._context_manager:
            self._context_manager.__exit__(exc_type, exc_val, exc_tb)
