"""Span enrichment helpers.

These functions set Langfuse-specific OTel span attributes so that
the Langfuse SDK recognises them when it exports traces.
"""

from __future__ import annotations

from typing import Any

from fastmcp.server.dependencies import get_context
from opentelemetry import trace as otel_trace
from opentelemetry.trace import format_trace_id

from .config import LangfuseConfig


def _tag_span_io(span: Any, *, input: str, output: str | None = None) -> None:
    """Set Langfuse input/output on both the trace AND observation level."""
    span.set_attribute("langfuse.trace.input", input)
    span.set_attribute("langfuse.observation.input", input)
    if output is not None:
        span.set_attribute("langfuse.trace.output", output)
        span.set_attribute("langfuse.observation.output", output)


def _enrich_current_span(config: LangfuseConfig) -> None:
    """Enrich the active OTel span with Langfuse session/env/trace-URL.

    FastMCP creates the OTel span *inside* the handler -- so this must be
    called from within a handler, not from middleware.

    Args:
        config: The active ``LangfuseConfig`` controlling which attributes
                are set on the span.
    """
    span = otel_trace.get_current_span()
    sc = span.get_span_context()
    if not sc or not sc.is_valid:
        return

    # Environment & release
    if config.environment:
        span.set_attribute("langfuse.environment", config.environment)
    if config.release:
        span.set_attribute("langfuse.release", config.release)

    # Session & user tracking
    if config.session_tracking:
        try:
            ctx = get_context()
            span.set_attribute("langfuse.session.id", ctx.session_id or "unknown")
            span.set_attribute("langfuse.user.id", ctx.client_id or "anonymous")
            span.set_attribute(
                "langfuse.trace.metadata.transport",
                str(ctx.transport or "unknown"),
            )
            span.set_attribute(
                "langfuse.trace.metadata.client_id",
                ctx.client_id or "anonymous",
            )
        except RuntimeError:
            pass

    if config.trace_urls:
        trace_id = format_trace_id(sc.trace_id)
        url = f"{config.base_url}/trace/{trace_id}"
        print(f"   🔗 Langfuse trace: {url}", flush=True)
