"""
Context management for mongoengine read-write router.

Uses Python's contextvars for thread-safe, async-compatible request context.
"""

from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any


@dataclass
class RequestContextData:
    """Request-scoped context data for read-write routing decisions."""

    user_id: str | None = None
    request_id: str | None = None
    write_operations: bool = False
    current_db_alias: str | None = None
    in_transaction: bool = False
    extra: dict[str, Any] = field(default_factory=dict)

    def mark_write(self) -> None:
        """Mark that a write operation occurred."""
        self.write_operations = True

    def set_db_alias(self, alias: str) -> None:
        """Set the current database alias being used."""
        self.current_db_alias = alias


_context_var: ContextVar[RequestContextData | None] = ContextVar(
    "mongoengine_rw_router_context", default=None
)


class RequestContext:
    """Thread-safe context manager for request-scoped data."""

    @staticmethod
    def set_context(
        user_id: str | None = None,
        request_id: str | None = None,
        extra: dict[str, Any] | None = None,
    ) -> RequestContextData:
        """
        Set a new context for the current request.

        Args:
            user_id: Optional user identifier for consistent hashing
            request_id: Optional request identifier for tracing
            extra: Optional extra data

        Returns:
            The created RequestContextData instance
        """
        ctx = RequestContextData(
            user_id=user_id,
            request_id=request_id,
            extra=extra or {},
        )
        _context_var.set(ctx)
        return ctx

    @staticmethod
    def get_context() -> RequestContextData | None:
        """Get the current context, or None if not set."""
        return _context_var.get(None)

    @staticmethod
    def reset() -> None:
        """Clear the current context."""
        _context_var.set(None)

    @staticmethod
    def update_context(**kwargs: Any) -> RequestContextData:
        """
        Update the current context with new values.

        Args:
            **kwargs: Fields to update (user_id, request_id, extra, etc.)

        Returns:
            The updated RequestContextData instance

        Raises:
            RuntimeError: If no context is set
        """
        ctx = _context_var.get(None)
        if ctx is None:
            ctx = RequestContextData()
            _context_var.set(ctx)

        for key, value in kwargs.items():
            if hasattr(ctx, key):
                setattr(ctx, key, value)
            else:
                ctx.extra[key] = value

        return ctx

    @staticmethod
    def mark_write() -> None:
        """Mark that a write operation occurred in current context."""
        ctx = _context_var.get(None)
        if ctx:
            ctx.mark_write()

    @staticmethod
    def set_in_transaction(in_transaction: bool) -> None:
        """Set whether we're currently in a transaction."""
        ctx = _context_var.get(None)
        if ctx:
            ctx.in_transaction = in_transaction


def get_user_id() -> str | None:
    """Get the current user_id from context."""
    ctx = _context_var.get(None)
    return ctx.user_id if ctx else None


def get_request_id() -> str | None:
    """Get the current request_id from context."""
    ctx = _context_var.get(None)
    return ctx.request_id if ctx else None


def has_write_operations() -> bool:
    """Check if any write operations occurred in current context."""
    ctx = _context_var.get(None)
    return ctx.write_operations if ctx else False


def is_in_transaction() -> bool:
    """Check if currently in a transaction."""
    ctx = _context_var.get(None)
    return ctx.in_transaction if ctx else False
