"""
Auth middleware for setting baggage context during authentication.

This middleware should be added BEFORE bind_otel_context in your FastAPI app.
"""

from typing import Optional
from fastapi import Request, HTTPException
from opentelemetry import context
from guardianhub import get_logger
from guardianhub.observability.context_utils import set_baggage_context

logger = get_logger(__name__)


async def set_auth_baggage_context(request: Request, call_next):
    """
    Auth middleware that extracts authentication info and sets baggage context.
    
    This MUST run before bind_otel_context to ensure baggage is available
    for extraction and propagation.
    """
    # Skip baggage setting for health/metrics endpoints
    if request.url.path in ["/health", "/metrics", "/docs", "/openapi.json"]:
        return await call_next(request)
    
    # Extract Authorization header
    auth_header = request.headers.get("Authorization")
    if not auth_header:
        logger.debug(
            "No Authorization header found, skipping baggage context",
            extra={"path": request.url.path}
        )
        return await call_next(request)
    
    try:
        # TODO: Replace this with your actual auth logic
        # This is where you would:
        # 1. Validate JWT token
        # 2. Look up user from database
        # 3. Get tenant information
        # 4. Extract session data
        
        # Extract tenant_id from prefixed token format: "Bearer tenant_id:actual_token"
        extracted_tenant_id = None
        if auth_header and ":" in auth_header:
            parts = auth_header.replace("Bearer ", "").split(":")
            if len(parts) > 1:
                extracted_tenant_id = parts[0]
        
        # MOCK IMPLEMENTATION - Replace with your real auth
        user_id = "user@example.com"  # Extract from validated token
        tenant_id = extracted_tenant_id or "tenant-123"  # Use extracted or fallback
        session_id = "session-456"     # Get from session store
        
        # Extract mission_id if present in query/path
        mission_id = request.query_params.get("mission_id") or \
                    request.path_params.get("mission_id")
        
        # Build baggage items dictionary
        baggage_items = {
            "user_id": user_id,
            "tenant_id": tenant_id,
            "session_id": session_id,
            "access_token": auth_header,  # Full token for downstream
        }
        
        # Add mission_id if present
        if mission_id:
            baggage_items["mission_id"] = mission_id
        
        # Use context_utils to properly set and attach baggage
        token = set_baggage_context(baggage_items)
        
        # Store token for cleanup
        request.state.otel_token = token
        
        logger.debug(
            "Auth baggage context set successfully",
            extra={
                "user_id": user_id,
                "tenant_id": tenant_id,
                "session_id": session_id,
                "has_mission_id": bool(mission_id),
                "path": request.url.path,
                "baggage_items": list(baggage_items.keys())
            }
        )
        
    except Exception as e:
        logger.error(
            "Failed to set auth baggage context",
            extra={
                "error": str(e),
                "path": request.url.path,
                "has_auth_header": bool(auth_header)
            }
        )
        # Continue without baggage context rather than failing the request
    
    # Continue to next middleware (bind_otel_context)
    try:
        return await call_next(request)
    finally:
        # Clean up the context token if it was set
        if hasattr(request.state, "otel_token"):
            logger.debug("set_auth_baggage_context: Detaching auth context token")
            context.detach(request.state.otel_token)


# Alternative: FastAPI dependency for specific endpoints
async def get_current_user_with_baggage(request: Request) -> dict:
    """
    Dependency that authenticates user and sets baggage context.
    
    Usage:
    @app.get("/api/protected-endpoint")
    async def protected_endpoint(current_user: dict = Depends(get_current_user_with_baggage)):
        # Your endpoint logic - baggage context is now available!
        pass
    """
    auth_header = request.headers.get("Authorization")
    if not auth_header:
        raise HTTPException(status_code=401, detail="No authorization header")
    
    # Your actual auth logic here
    # Validate token, get user info, etc.
    user_id = "user@example.com"
    tenant_id = "tenant-123"
    session_id = "session-456"
    
    # Build baggage items dictionary
    baggage_items = {
        "user_id": user_id,
        "tenant_id": tenant_id,
        "session_id": session_id,
        "access_token": auth_header,
    }
    
    # Use context_utils to set baggage context
    token = set_baggage_context(baggage_items)
    request.state.otel_token = token
    
    logger.info(
        "Baggage context set via dependency",
        extra={
            "user_id": user_id,
            "tenant_id": tenant_id,
            "path": request.url.path
        }
    )
    
    return {"user_id": user_id, "tenant_id": tenant_id, "session_id": session_id}
