from typing import Optional

from fastapi import Request
from opentelemetry import baggage
from opentelemetry import trace
from opentelemetry.context import Context
from opentelemetry.propagate import extract
from opentelemetry.sdk.trace import SpanProcessor
from opentelemetry.sdk.trace.sampling import (
    Sampler,
    SamplingResult,
    Decision,
)
from opentelemetry.trace import Span


async def bind_otel_context(request: Request, call_next):
    # 1️⃣ Extract OTEL context from headers
    ctx = extract(request.headers)

    # 2️⃣ Attach context so downstream sees it
    with trace.use_span(trace.get_current_span(ctx), end_on_exit=False):
        return await call_next(request)


class GuardianHubSampler(Sampler):
    """
    Drops low-level ASGI spans (http receive/send)
    Keeps route, agent, and custom spans.
    """

    def should_sample(
        self,
        parent_context,
        trace_id,
        name,
        kind=None,
        attributes=None,
        links=None,
    ) -> SamplingResult:

        lname = (name or "").lower()

        if "http receive" in lname or "http send" in lname:
            return SamplingResult(Decision.DROP)

        return SamplingResult(Decision.RECORD_AND_SAMPLE)

    def get_description(self) -> str:
        return "GuardianHub Sampler (drop ASGI receive/send)"

class BaggageToSpanProcessor(SpanProcessor):
    """
    Copies selected baggage values into span attributes
    so they appear on every span.
    """

    BAGGAGE_KEYS = (
        "user_id",
        "session_id",
    )

    def on_start( self,
        span: Span,
        parent_context: Optional[Context] = None,
    ) -> None:
        for key in self.BAGGAGE_KEYS:
            value = baggage.get_baggage(key, parent_context)
            if value is not None:
                span.set_attribute(key, str(value))

    def on_end(self, span):
        pass