from typing import Optional

from fastapi import Request
from opentelemetry import baggage, propagate, context
from opentelemetry.context import Context
from opentelemetry.sdk.trace import SpanProcessor
from opentelemetry.trace import Span

from guardianhub import get_logger

logger = get_logger(__name__)


async def bind_otel_context(request: Request, call_next):
    logger.debug(
        "bind_otel_context: Extracting context from request headers",
        extra={
            "method": request.method,
            "url": str(request.url),
            "headers": dict(request.headers)
        }
    )
    
    ctx = propagate.extract(request.headers)
    
    # Log extracted baggage for debugging
    extracted_baggage = {}
    try:
        # Try to get common baggage keys to show what was extracted
        for key in ["user_id", "tenant_id", "session_id", "access_token", "mission_id"]:
            value = baggage.get_baggage(key, ctx)
            if value:
                extracted_baggage[key] = str(value)[:50] + "..." if len(str(value)) > 50 else str(value)
    except Exception as e:
        logger.warning(
            "bind_otel_context: Error extracting baggage for logging",
            extra={"error": str(e)}
        )
    
    logger.debug(
        "bind_otel_context: Context extracted",
        extra={
            "extracted_baggage": extracted_baggage,
            "method": request.method,
            "url": str(request.url)
        }
    )
    
    token = context.attach(ctx)
    
    try:
        logger.debug(
            "bind_otel_context: Context attached, calling next middleware",
            extra={
                "method": request.method,
                "url": str(request.url)
            }
        )
        
        response = await call_next(request)
        
        logger.debug(
            "bind_otel_context: Next middleware completed",
            extra={
                "method": request.method,
                "url": str(request.url),
                "status_code": response.status_code
            }
        )
        
        return response
    finally:
        # DETACH the original extraction
        logger.debug("bind_otel_context: Detaching original context token")
        context.detach(token)
        
        # DETACH the internal Auth packing if it exists
        if hasattr(request.state, "otel_token"):
            logger.debug("bind_otel_context: Detaching auth context token")
            context.detach(request.state.otel_token)
        else:
            logger.debug("bind_otel_context: No auth context token to detach")




# Example Auth Middleware / Dependency - implement as needed for your auth system
# async def get_current_user_and_pack_context(request: Request, db_session: Session):
#     # 1. Your existing Auth/DB logic
#     user_token = request.headers.get("Authorization")
#     tenant_data = await db_session.execute(...)  # Fetch your TenantPacket
#
#     # 2. THE 3 LINES: Pack the Sovereign Context
#     # This 'activates' the data so the SovereignContextProcessor can see it
#     ctx = baggage.set_baggage("tenant_id", tenant_data.id)
#     ctx = baggage.set_baggage("user_id", tenant_data.user_id)
#     ctx = baggage.set_baggage("access_token", user_token)  # For downstream mesh tools
#
#     # 3. Attach it to the current thread
#     token = context.attach(ctx)
#
#     # Crucial: Store the token in request.state so we can detach it later
#     request.state.otel_token = token
#
#     return tenant_data


class SovereignContextProcessor(SpanProcessor):
    """
    Ensures Identity, Auth, and Tenant context are visible to every
    service, database, and tool in the mesh.
    """
    # Keys we want to propagate AND stamp on spans
    PROPAGATION_KEYS = ("user_id", "tenant_id", "access_token", "mission_id")

    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
        logger.debug(
            "SovereignContextProcessor: Starting span processing",
            extra={
                "span_name": span.name,
                "span_kind": span.kind.name if span.kind else "unknown",
                "trace_id": span.get_span_context().trace_id,
                "span_id": span.get_span_context().span_id
            }
        )
        
        processed_keys = []
        
        for key in self.PROPAGATION_KEYS:
            value = baggage.get_baggage(key, parent_context)
            if value:
                # Stamp for Observability (Langfuse/Tempo)
                span.set_attribute(f"ctx.{key}", str(value))
                processed_keys.append(key)
                logger.debug(
                    f"SovereignContextProcessor: Set baggage attribute {key}",
                    extra={
                        "key": key,
                        "value": str(value)[:50] + "..." if len(str(value)) > 50 else str(value),
                        "span_name": span.name
                    }
                )

                # Special Mapping for Langfuse native fields
                session_id = baggage.get_baggage("session_id", parent_context)
                if session_id:
                    span.set_attribute("langfuse.session.id", str(session_id))
                    span.set_attribute("ctx.session_id", str(session_id))
                    logger.debug(
                        "SovereignContextProcessor: Set Langfuse session ID",
                        extra={
                            "session_id": str(session_id),
                            "span_name": span.name
                        }
                    )
                    
                if key == "user_id":
                    span.set_attribute("langfuse.user.id", str(value))
                    logger.debug(
                        "SovereignContextProcessor: Set Langfuse user ID",
                        extra={
                            "user_id": str(value),
                            "span_name": span.name
                        }
                    )
                    
                if key == "tenant_id":
                    span.set_attribute("langfuse.trace.metadata.tenant", str(value))
                    logger.debug(
                        "SovereignContextProcessor: Set Langfuse tenant metadata",
                        extra={
                            "tenant_id": str(value),
                            "span_name": span.name
                        }
                    )
            else:
                logger.debug(
                    f"SovereignContextProcessor: No baggage value for key {key}",
                    extra={
                        "key": key,
                        "span_name": span.name
                    }
                )
        
        if processed_keys:
            logger.info(
                "SovereignContextProcessor: Successfully processed context baggage",
                extra={
                    "processed_keys": processed_keys,
                    "span_name": span.name,
                    "trace_id": span.get_span_context().trace_id
                }
                )
        else:
            logger.debug(
                "SovereignContextProcessor: No baggage keys found for processing",
                extra={
                    "span_name": span.name,
                    "trace_id": span.get_span_context().trace_id
                }
            )

    def on_end(self, span: Span) -> None:
        logger.debug(
            "SovereignContextProcessor: Ending span processing",
            extra={
                "span_name": span.name,
                "trace_id": span.get_span_context().trace_id,
                "span_id": span.get_span_context().span_id,
                "status": span.status.status_code.name if span.status else "unknown"
            }
        )
        pass