"""Request correlation tracking using contextvars.

Provides correlation ID tracking for linking related operations
across HTTP requests, database queries, and other I/O.
"""

from __future__ import annotations

import uuid
from contextvars import ContextVar
from typing import Optional, Any

# Try to use asgi-correlation-id if available
try:
    from asgi_correlation_id import CorrelationIdMiddleware
    from asgi_correlation_id.context import correlation_id as asgi_correlation_id

    ASGI_CORRELATION_AVAILABLE = True
except ImportError:
    ASGI_CORRELATION_AVAILABLE = False
    asgi_correlation_id = None
    CorrelationIdMiddleware = None

# Fallback correlation ID for non-ASGI contexts
_fallback_correlation_id: ContextVar[Optional[str]] = ContextVar(
    "fallback_correlation_id", default=None
)

# Additional context for request tracking
_request_context: ContextVar[Optional[dict]] = ContextVar("request_context", default=None)


def get_correlation_id() -> str:
    """Get current correlation ID, generating if needed.

    Returns:
        8-character correlation ID
    """
    # Try asgi-correlation-id first
    if ASGI_CORRELATION_AVAILABLE and asgi_correlation_id:
        cid = asgi_correlation_id.get()
        if cid:
            return cid

    # Check fallback
    cid = _fallback_correlation_id.get()
    if not cid:
        cid = str(uuid.uuid4())[:8]
        _fallback_correlation_id.set(cid)

    return cid


def get_async_task_id() -> Optional[str]:
    """Return current asyncio task id if available."""
    try:
        import asyncio
        task = asyncio.current_task()
        if task is None:
            return None
        return str(id(task))
    except Exception:
        return None


def set_correlation_id(cid: str) -> None:
    """Set correlation ID (for non-ASGI contexts).

    Args:
        cid: Correlation ID to set
    """
    _fallback_correlation_id.set(cid)


def clear_correlation_id() -> None:
    """Clear the current correlation ID."""
    _fallback_correlation_id.set(None)


def get_correlation_middleware():
    """Get the appropriate correlation middleware for ASGI apps.

    Returns:
        CorrelationIdMiddleware class

    Raises:
        ImportError: If asgi-correlation-id is not installed
    """
    if not ASGI_CORRELATION_AVAILABLE or CorrelationIdMiddleware is None:
        raise ImportError(
            "asgi-correlation-id not installed. "
            "Install with: pip install asgi-correlation-id"
        )

    return CorrelationIdMiddleware


class CorrelationContext:
    """Context manager for correlation ID in sync code.

    Usage:
        with CorrelationContext("my-id") as cid:
            # All operations here have correlation ID "my-id"
            pass

        # Generates new ID if not provided
        with CorrelationContext() as cid:
            print(f"Using correlation ID: {cid}")
    """

    def __init__(self, correlation_id: Optional[str] = None):
        """Initialize with optional correlation ID.

        Args:
            correlation_id: ID to use, or None to generate
        """
        self.correlation_id = correlation_id or str(uuid.uuid4())[:8]
        self._token = None

    def __enter__(self) -> str:
        """Enter context, setting correlation ID."""
        self._token = _fallback_correlation_id.set(self.correlation_id)
        return self.correlation_id

    def __exit__(self, *args: Any) -> None:
        """Exit context, restoring previous correlation ID."""
        if self._token is not None:
            _fallback_correlation_id.reset(self._token)


class AsyncCorrelationContext:
    """Async context manager for correlation ID.

    Usage:
        async with AsyncCorrelationContext("my-id") as cid:
            await some_async_operation()
    """

    def __init__(self, correlation_id: Optional[str] = None):
        """Initialize with optional correlation ID.

        Args:
            correlation_id: ID to use, or None to generate
        """
        self.correlation_id = correlation_id or str(uuid.uuid4())[:8]
        self._token = None

    async def __aenter__(self) -> str:
        """Async enter context."""
        self._token = _fallback_correlation_id.set(self.correlation_id)
        return self.correlation_id

    async def __aexit__(self, *args: Any) -> None:
        """Async exit context."""
        if self._token is not None:
            _fallback_correlation_id.reset(self._token)


def set_request_context(**kwargs: Any) -> None:
    """Set additional request context.

    Args:
        **kwargs: Key-value pairs to store in context
    """
    ctx = _request_context.get()
    if ctx is None:
        ctx = {}
    ctx.update(kwargs)
    _request_context.set(ctx)


def get_request_context() -> dict:
    """Get current request context.

    Returns:
        Dictionary of request context values
    """
    return _request_context.get() or {}


def clear_request_context() -> None:
    """Clear the request context."""
    _request_context.set({})


def correlation_decorator(func):
    """Decorator to automatically add correlation ID to function calls.

    Usage:
        @correlation_decorator
        def my_function():
            # Automatically has correlation ID
            pass
    """

    def wrapper(*args, **kwargs):
        with CorrelationContext():
            return func(*args, **kwargs)

    return wrapper


def async_correlation_decorator(func):
    """Async decorator to automatically add correlation ID.

    Usage:
        @async_correlation_decorator
        async def my_async_function():
            # Automatically has correlation ID
            pass
    """

    async def wrapper(*args, **kwargs):
        async with AsyncCorrelationContext():
            return await func(*args, **kwargs)

    return wrapper
