from typing import Optional

from fastapi import Request
from opentelemetry import baggage
from opentelemetry import propagate, context
from opentelemetry.context import Context
from opentelemetry.sdk.trace import SpanProcessor
from opentelemetry.trace import Span
from opentelemetry import baggage, context
from sqlalchemy.orm import Session  # <--- Add this
from opentelemetry import baggage, context


async def bind_otel_context(request: Request, call_next):
    ctx = propagate.extract(request.headers)
    token = context.attach(ctx)
    try:
        response = await call_next(request)
        return response
    finally:
        # DETACH the original extraction
        context.detach(token)
        # DETACH the internal Auth packing if it exists
        if hasattr(request.state, "otel_token"):
            context.detach(request.state.otel_token)




# Inside your Auth Middleware / Dependency
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:
        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))

                # 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))
                if key == "user_id":
                    span.set_attribute("langfuse.user.id", str(value))
                if key == "tenant_id":
                    span.set_attribute("langfuse.trace.metadata.tenant", str(value))

    def on_end(self, span: Span) -> None:
        pass