"""Request-scoped runtime state."""

from contextlib import AbstractContextManager, contextmanager
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from typing import Iterator
from uuid import uuid4

from punq import Container  # type: ignore[import-untyped]

from ..support.authz import Principal
from ..support.errors import MissingRequestScopeError


@dataclass(slots=True, frozen=True)
class RequestState:
    correlation_id: str = "unknown"
    principal: Principal = field(default_factory=Principal.guest)
    container: Container | None = None


_request_state: ContextVar[RequestState] = ContextVar("request_state", default=RequestState())


class RequestScope(AbstractContextManager[RequestState]):
    """Activate request state for the current invocation."""

    def __init__(
        self,
        *,
        correlation_id: str | None = None,
        principal: Principal | None = None,
        container: Container | None = None,
    ) -> None:
        self._correlation_id = correlation_id
        self._principal = principal
        self._container = container
        self._token: Token[RequestState] | None = None

    def __enter__(self) -> RequestState:
        current = _request_state.get()
        next_state = RequestState(
            correlation_id=self._correlation_id or _next_correlation_id(current),
            principal=self._principal or current.principal,
            container=self._container if self._container is not None else current.container,
        )
        self._token = _request_state.set(next_state)
        return next_state

    def __exit__(self, exc_type, exc, tb) -> None:
        if self._token is not None:
            _request_state.reset(self._token)
            self._token = None
        return None


def _next_correlation_id(current: RequestState) -> str:
    if current.correlation_id != "unknown":
        return current.correlation_id
    return str(uuid4())


def current_request_state() -> RequestState:
    """Return the active request state."""

    return _request_state.get()


def current_container() -> Container | None:
    """Return the active request container when available."""

    return current_request_state().container


def get_app_context() -> Container:
    """Return the active request container."""

    container = current_container()
    if container is None:
        raise MissingRequestScopeError("No active request container. Use RequestScope or build_lambda_handler().")
    return container


def correlation_id() -> str:
    """Return the active correlation id."""

    return current_request_state().correlation_id


def principal() -> Principal:
    """Return the active principal."""

    return current_request_state().principal


def set_principal(value: Principal) -> None:
    """Replace the principal in the current request state."""

    state = current_request_state()
    _request_state.set(RequestState(correlation_id=state.correlation_id, principal=value, container=state.container))


@contextmanager
def use_correlation_id(value: str) -> Iterator[None]:
    """Override the correlation id for the active scope."""

    state = current_request_state()
    token = _request_state.set(RequestState(correlation_id=value, principal=state.principal, container=state.container))
    try:
        yield
    finally:
        _request_state.reset(token)


__all__ = [
    "RequestScope",
    "RequestState",
    "correlation_id",
    "current_container",
    "current_request_state",
    "get_app_context",
    "principal",
    "set_principal",
    "use_correlation_id",
]
