"""Context manager utilities for manual telemetry spans."""

from __future__ import annotations

import json
import logging
import os
import random
from collections.abc import AsyncGenerator, Generator, Mapping, Sequence
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final

if TYPE_CHECKING:
    from basalt.experiments import Experiment
    from basalt.prompts.models import Prompt

from opentelemetry import context as otel_context
from opentelemetry import trace
from opentelemetry.context import attach, detach, set_value
from opentelemetry.trace import Span, Tracer

from ..types.common import JSONValue, SpanAttributeValue
from . import semconv
from .trace_context import (
    ORGANIZATION_CONTEXT_KEY,
    SHOULD_EVALUATE_CONTEXT_KEY,
    USER_CONTEXT_KEY,
    TraceIdentity,
    _current_trace_defaults,
    _TraceContextConfig,
    apply_organization_from_context,
    apply_user_from_context,
)
from .types import Identity

SPAN_TYPE_ATTRIBUTE = semconv.BasaltSpan.KIND
EVALUATOR_CONTEXT_KEY: Final[str] = "basalt.context.evaluators"
EVALUATOR_CONFIG_CONTEXT_KEY: Final[str] = "basalt.context.evaluator_config"
EVALUATOR_METADATA_CONTEXT_KEY: Final[str] = "basalt.context.evaluator_metadata"
ROOT_SPAN_CONTEXT_KEY: Final[str] = "basalt.context.root_span"
ROOT_PARENT_SPAN_CONTEXT_KEY: Final[str] = "basalt.context.root_parent_span"
logger = logging.getLogger(__name__)


class NoOpSpanHandle:
    """No-op span handle that does nothing.

    Used when observe() is called without a preceding start_observe.
    This prevents creating Basalt spans when not in a trace context.
    """

    def set_attribute(self, key: str, value: str | int | float | bool | None) -> None:
        """No-op: does nothing."""
        pass

    def set_attributes(self, attributes: dict[str, Any]) -> None:
        """No-op: does nothing."""
        pass

    def set_metadata(self, metadata: Mapping[str, Any] | None) -> None:
        """No-op: does nothing."""
        pass

    def set_prompt(self, prompt: Any) -> None:
        """No-op: does nothing."""
        pass

    def set_input(self, payload: JSONValue) -> None:
        """No-op: does nothing."""
        pass

    def set_output(self, payload: JSONValue) -> None:
        """No-op: does nothing."""
        pass

    def set_io(
        self,
        *,
        input_payload: JSONValue | None = None,
        output_payload: JSONValue | None = None,
        variables: Mapping[str, Any] | None = None,
    ) -> None:
        """No-op: does nothing."""
        pass

    def add_evaluator(self, evaluator_slug: str) -> None:
        """No-op: does nothing."""
        pass

    def add_evaluators(self, *evaluators: Sequence[str]) -> None:
        """No-op: does nothing."""
        pass

    def set_identity(self, identity: Mapping[str, Any] | None = None) -> None:
        """No-op: does nothing."""
        pass

    def set_model(self, model: str) -> None:
        """No-op: does nothing."""
        pass

    def set_response_model(self, model: str) -> None:
        """No-op: does nothing."""
        pass

    def set_operation_name(self, operation: str) -> None:
        """No-op: does nothing."""
        pass

    def __enter__(self):
        """No-op: does nothing."""
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """No-op: does nothing."""
        pass


@dataclass(slots=True, frozen=True)
class EvaluationConfig:
    """
    Type-safe configuration for evaluators attached to a span.

    This configuration is span-scoped and controls trace-level sampling for evaluators.
    The sample_rate determines whether evaluators run for the entire trace.

    Attributes:
        sample_rate: Sampling rate for trace-level evaluation (0.0-1.0). Default is 0.0 (no sampling).
                     When set, one sampling decision is made at root span creation and propagated
                     to all spans in the trace via basalt.span.should_evaluate attribute.
    """

    sample_rate: float = 0.0

    def __post_init__(self) -> None:
        if not 0.0 <= self.sample_rate <= 1.0:
            raise ValueError("sample_rate must be within [0.0, 1.0].")

    def to_dict(self) -> dict[str, Any]:
        """Convert config to a dictionary for serialization."""
        return {"sample_rate": self.sample_rate}


@dataclass(slots=True)
class EvaluatorAttachment:
    """Normalized evaluator payload applied to spans."""

    slug: str
    metadata: Mapping[str, Any] | None = field(default=None)

    def __post_init__(self) -> None:
        if not isinstance(self.slug, str) or not self.slug.strip():
            raise ValueError("Evaluator slug must be a non-empty string.")
        self.slug = self.slug.strip()
        if self.metadata is not None and not isinstance(self.metadata, Mapping):
            raise TypeError("Evaluator metadata must be a mapping.")


def _normalize_evaluator_entry(entry: object) -> EvaluatorAttachment:
    """Convert assorted evaluator payloads into EvaluatorAttachment objects."""
    if isinstance(entry, EvaluatorAttachment):
        return entry
    if isinstance(entry, str):
        return EvaluatorAttachment(slug=entry)
    if isinstance(entry, Mapping):
        payload = dict(entry)
        slug = payload.pop("slug", None)
        if slug is None:
            raise ValueError("Evaluator mapping must include a 'slug' key.")
        metadata = payload.pop("metadata", None)
        return EvaluatorAttachment(slug=str(slug), metadata=metadata)
    raise TypeError(f"Unsupported evaluator specification: {entry!r}")


def _normalize_evaluators(evaluators: Sequence[Any] | None) -> list[EvaluatorAttachment]:
    if not evaluators:
        return []
    if not isinstance(evaluators, Sequence) or isinstance(evaluators, (str, bytes)):
        evaluators = [evaluators]
    result: list[EvaluatorAttachment] = []
    for entry in evaluators:
        if isinstance(entry, Sequence) and not isinstance(entry, (str, bytes, EvaluatorAttachment)):
            result.extend(_normalize_evaluators(entry))
        else:
            result.append(_normalize_evaluator_entry(entry))
    return result


def normalize_evaluator_specs(evaluators: Sequence[Any] | None) -> list[EvaluatorAttachment]:
    """Public helper to normalize evaluator specifications."""
    return _normalize_evaluators(evaluators)


@contextmanager
def with_evaluators(
    evaluators: Sequence[Any],
) -> Generator[None, None, None]:
    """Propagate evaluator slugs, config, and metadata through the OpenTelemetry context.

    Args:
        evaluators: Evaluator specifications to propagate.

    """

    attachments = normalize_evaluator_specs(evaluators)
    # Only short-circuit when nothing at all provided. Empty metadata should still attach.

    # Propagate evaluator slugs
    tokens = []
    if attachments:
        existing = otel_context.get_value(EVALUATOR_CONTEXT_KEY)
        combined: list[str] = []
        if isinstance(existing, (list, tuple)):
            combined.extend(str(slug) for slug in existing if str(slug).strip())

        for attachment in attachments:
            if attachment.slug not in combined:
                combined.append(attachment.slug)

        tokens.append(attach(set_value(EVALUATOR_CONTEXT_KEY, tuple(combined))))

    try:
        yield
    finally:
        for token in reversed(tokens):
            detach(token)


def _attach_attributes(span: Span, attributes: dict[str, Any] | None) -> None:
    if not attributes:
        return
    for key, value in attributes.items():
        _set_serialized_attribute(span, key, value)


def _serialize_attribute(value: JSONValue) -> SpanAttributeValue:
    if value is None or isinstance(value, (str, bool, int, float)):
        return value
    try:
        return json.dumps(value)
    except Exception:
        return str(value)


def _set_serialized_attribute(span: Span, key: str, value: JSONValue) -> None:
    serialized = _serialize_attribute(value)
    if serialized is not None:
        span.set_attribute(key, serialized)


def trace_content_enabled() -> bool:
    flag = os.getenv("TRACELOOP_TRACE_CONTENT")
    if flag is None:
        return True
    return flag.strip().lower() not in {"0", "false", "no", "off"}


def get_tracer(tracer_name: str = "basalt.observability") -> Tracer:
    """Get or create a tracer with the given name."""
    return trace.get_tracer(tracer_name)


def get_current_otel_span() -> Span | None:  # Lightweight alias
    """Return the active OpenTelemetry span if valid, else None."""
    span = trace.get_current_span()
    if span is None or not span.get_span_context().is_valid:
        return None
    return span


def get_current_span_handle() -> SpanHandle | None:
    """Return a handle for the current span."""
    span = get_current_otel_span()
    if not span:
        return None
    return SpanHandle(span)


def get_root_span_handle() -> StartSpanHandle | None:
    """Return a handle for the root span of the current trace.

    This allows accessing the root span from deeply nested contexts,
    enabling late-binding of identify() or metadata operations.

    Returns:
        StartSpanHandle if a root span exists, None otherwise.
    """
    root_span = otel_context.get_value(ROOT_SPAN_CONTEXT_KEY)
    if root_span and isinstance(root_span, Span):
        parent_span = otel_context.get_value(ROOT_PARENT_SPAN_CONTEXT_KEY)
        if not isinstance(parent_span, Span):
            parent_span = None
        return StartSpanHandle(root_span, parent_span)
    return None


class SpanHandle:
    """Helper around an OTEL span with convenience methods."""

    def __init__(
        self,
        span: Span,
        parent_span: Span | None = None,
        defaults: _TraceContextConfig | None = None,
    ) -> None:
        self._span = span
        self._io_payload: dict[str, Any] = {"input": None, "output": None, "variables": None}
        self._parent_span = parent_span if parent_span and parent_span.get_span_context().is_valid else None
        self._evaluators: dict[str, EvaluatorAttachment] = {}
        self._evaluator_config: EvaluationConfig | None = None
        self._hydrate_existing_evaluators()

    def set_attribute(self, key: str, value: str | int | float | bool | None) -> None:
        """
        Sets metadata on the current span.

        Args:
            key (str): The metadata key to set.
            value (str | int | float | bool | None): The metadata value.

        Returns:
            None
        """
        if value is not None:
            self._span.set_attribute(key, value)

    def set_attributes(self, attributes: dict[str, Any]) -> None:
        """
        Set multiple attributes on the current span.

        Args:
            attributes: Dictionary of attributes to set. Complex values (dicts, lists)
                       will be serialized to JSON strings.
        """
        _attach_attributes(self._span, attributes)

    def set_metadata(self, metadata: Mapping[str, Any] | None) -> None:
        """Merge metadata onto span as a JSON object.

        Metadata is stored as a JSON-serialized dictionary at ``semconv.BasaltSpan.METADATA``.
        Existing metadata is preserved and shallow-merged with the incoming mapping.
        """
        if metadata is None:
            return
        if not isinstance(metadata, Mapping):
            raise TypeError("metadata must be a mapping")
        # Lazy import to avoid circular dependency
        from .utils import apply_span_metadata

        apply_span_metadata(self._span, metadata)

    def set_prompt(self, prompt: Prompt) -> None:
        """
        Set prompt metadata on the current span.

        Args:
            prompt: The Prompt object to attach.
        """
        # Prepare prompt metadata for span attributes
        prompt_metadata = {
            "basalt.prompt.slug": prompt.slug,
            "basalt.prompt.version": prompt.version,
            "basalt.prompt.model.provider": prompt.model.provider,
            "basalt.prompt.model.model": prompt.model.model,
        }

        # Store prompt.variables separately if available
        if prompt.variables:
            prompt_metadata["basalt.prompt.variables"] = json.dumps(prompt.variables)

        self.set_attributes(prompt_metadata)

    # ------------------------------------------------------------------
    # IO helpers
    # ------------------------------------------------------------------
    def set_input(self, payload: JSONValue) -> None:
        """
        Sets the input payload for the current context manager.
        Stores the provided payload in the internal `_io_payload` dictionary under the "input" key.
        If trace content is enabled, serializes and attaches the input payload to the tracing span
        using the appropriate semantic convention.
        Args:
            payload (str): The input data to be recorded and optionally traced.
        """

        self._io_payload["input"] = payload
        if trace_content_enabled():
            _set_serialized_attribute(self._span, semconv.BasaltSpan.INPUT, payload)

    def set_output(self, payload: JSONValue) -> None:
        """
        Sets the output payload for the current context manager.
        Stores the provided payload in the internal I/O payload dictionary under the "output" key.
        If trace content is enabled, serializes and attaches the payload to the current span for observability.
        Args:
            payload (Any): The output data to be stored and optionally traced.
        """

        self._io_payload["output"] = payload
        if trace_content_enabled():
            _set_serialized_attribute(self._span, semconv.BasaltSpan.OUTPUT, payload)

    def set_io(
        self,
        *,
        input_payload: JSONValue | None = None,
        output_payload: JSONValue | None = None,
        variables: Mapping[str, Any] | None = None,
    ) -> None:
        """
        Sets the input, output, and variables payloads for the current context manager.
        """
        if input_payload is not None:
            self.set_input(input_payload)
        if output_payload is not None:
            self.set_output(output_payload)
        if variables is not None:
            if not isinstance(variables, Mapping):
                raise TypeError("Span variables must be provided as a mapping.")
            self._io_payload["variables"] = dict(variables)
            if trace_content_enabled():
                _set_serialized_attribute(self._span, semconv.BasaltSpan.VARIABLES, dict(variables))
                if self._parent_span:
                    _set_serialized_attribute(self._parent_span, semconv.BasaltSpan.VARIABLES, dict(variables))

    def _io_snapshot(self) -> dict[str, Any]:
        """Return a shallow copy of the tracked IO payload."""
        snapshot = dict(self._io_payload)
        if snapshot["variables"] is not None:
            snapshot["variables"] = dict(snapshot["variables"])
        return snapshot

    def add_evaluator(
        self,
        evaluator_slug: str,
    ) -> None:
        """
        Attach an evaluator slug to the span.

        Args:
            evaluator_slug: The evaluator slug to attach.
        """
        attachment = EvaluatorAttachment(slug=evaluator_slug, metadata=None)
        self._append_evaluator(attachment)

    def add_evaluators(self, *evaluators: Sequence[str]) -> None:
        """Attach multiple evaluators to the span."""
        for attachment in normalize_evaluator_specs(evaluators):
            self._append_evaluator(attachment)

    def set_identity(self, identity: Identity | None = None) -> None:
        """
        Set user and/or organization identity for the span.

        Args:
            identity: Identity TypedDict with optional 'user' and 'organization' keys.
                Each key should contain a dict with 'id' (required) and 'name' (optional).

        Example:
            >>> span.set_identity(
            ...     {
            ...         "user": {"id": "user-123", "name": "John Doe"},
            ...         "organization": {"id": "org-456", "name": "ACME Corp"},
            ...     }
            ... )
        """
        if identity is None:
            return

        user_spec = identity.get("user") if identity else None
        org_spec = identity.get("organization") if identity else None

        if user_spec is not None:
            apply_user_from_context(self._span, user_spec)
        if org_spec is not None:
            apply_organization_from_context(self._span, org_spec)

    def _append_evaluator(self, attachment: EvaluatorAttachment) -> None:
        """Attach an evaluator to the span, avoiding duplicates.

        If the attachment includes metadata, it will be merged into the span-level metadata.
        """
        existing = self._evaluators.get(attachment.slug)
        if existing and existing == attachment:
            return
        self._evaluators[attachment.slug] = attachment
        evaluator_list = list(self._evaluators.keys())
        self._span.set_attribute(semconv.BasaltSpan.EVALUATORS, evaluator_list)

    def _hydrate_existing_evaluators(self) -> None:
        """Populate evaluator cache from span attributes if present."""

        attributes = getattr(self._span, "attributes", None)
        if not isinstance(attributes, dict):
            return
        current = attributes.get(semconv.BasaltSpan.EVALUATORS)
        if not isinstance(current, (list, tuple)):
            return
        for slug in current:
            if isinstance(slug, str) and slug.strip():
                normalized = EvaluatorAttachment(slug=slug)
                self._evaluators.setdefault(normalized.slug, normalized)

    @property
    def span(self) -> Span:
        return self._span

    # --- GenAI-specific convenience methods (promoted from LLMSpanHandle) ---

    def set_model(self, model: str) -> None:
        """Set the GenAI request model name."""
        self.set_attribute(semconv.GenAI.REQUEST_MODEL, model)

    def set_response_model(self, model: str) -> None:
        """Set the GenAI response model name."""
        self.set_attribute(semconv.GenAI.RESPONSE_MODEL, model)

    def set_operation_name(self, operation: str) -> None:
        """Set the GenAI operation name (e.g., 'chat', 'completion', 'embeddings')."""
        self.set_attribute(semconv.GenAI.OPERATION_NAME, operation)

    def set_provider(self, provider: str) -> None:
        """Set the GenAI provider name (e.g., 'openai', 'anthropic', 'google')."""
        self.set_attribute(semconv.GenAI.PROVIDER_NAME, provider)

    def set_tokens(
        self,
        *,
        input: int | None = None,
        output: int | None = None,
    ) -> None:
        """Set token usage counts for the GenAI operation."""
        if input is not None:
            self.set_attribute(semconv.GenAI.USAGE_INPUT_TOKENS, input)
        if output is not None:
            self.set_attribute(semconv.GenAI.USAGE_OUTPUT_TOKENS, output)

    def set_temperature(self, temperature: float) -> None:
        """Set the temperature parameter for the GenAI request."""
        self.set_attribute(semconv.GenAI.REQUEST_TEMPERATURE, temperature)

    def set_top_p(self, top_p: float) -> None:
        """Set the top_p parameter for the GenAI request."""
        self.set_attribute(semconv.GenAI.REQUEST_TOP_P, top_p)

    def set_top_k(self, top_k: float) -> None:
        """Set the top_k parameter for the GenAI request."""
        self.set_attribute(semconv.GenAI.REQUEST_TOP_K, top_k)

    def set_max_tokens(self, max_tokens: int) -> None:
        """Set the max_tokens parameter for the GenAI request."""
        self.set_attribute(semconv.GenAI.REQUEST_MAX_TOKENS, max_tokens)

    def set_response_id(self, response_id: str) -> None:
        """Set the GenAI response ID (completion ID)."""
        self.set_attribute(semconv.GenAI.RESPONSE_ID, response_id)

    def set_finish_reasons(self, reasons: list[str]) -> None:
        """Set the finish reasons array for the GenAI response."""
        _set_serialized_attribute(self._span, semconv.GenAI.RESPONSE_FINISH_REASONS, list(reasons))

    def set_latency(self, duration_ms: float) -> None:
        """Set the GenAI operation latency in milliseconds."""
        self.set_attribute(semconv.GenAI.LATENCY, duration_ms)


class StartSpanHandle(SpanHandle):
    """
    Span handle for root spans created by StartObserve.

    This handle type supports experiment and evaluation configuration,
    which are only applicable to root spans in a trace.
    """

    def set_evaluation_config(self, config: EvaluationConfig | Mapping[str, Any]) -> None:
        """Attach span-scoped evaluator configuration.

        The configuration applies to all evaluators attached to this span.
        It is stored under the semantic key BasaltSpan.EVALUATION_CONFIG as JSON.

        Args:
            config: Either an EvaluationConfig instance or a mapping with config values.
        """
        if isinstance(config, EvaluationConfig):
            self._evaluator_config = config
            _set_serialized_attribute(self._span, semconv.BasaltSpan.EVALUATION_CONFIG, config.to_dict())
        elif isinstance(config, Mapping):
            config_dict = dict(config)
            self._evaluator_config = EvaluationConfig(**config_dict)
            _set_serialized_attribute(self._span, semconv.BasaltSpan.EVALUATION_CONFIG, config_dict)
        else:
            raise TypeError("Evaluator config must be an EvaluationConfig or a mapping.")

    def set_experiment(
        self,
        experiment: str | Experiment,
    ) -> None:
        """
        Set the experiment identity for the root span.

        This method is only available on root spans (StartSpanHandle).
        """
        # Use duck typing instead of isinstance to avoid circular import
        experiment_id: str
        if hasattr(experiment, "id"):
            experiment_id = experiment.id  # type: ignore[attr-defined]
        else:
            experiment_id = str(experiment)

        self._span.set_attribute(semconv.BasaltExperiment.ID, experiment_id)


class LLMSpanHandle(SpanHandle):
    """
    Span handle for LLM/GenAI operations.

    Provides specialized methods for LLM operations that require structured message formatting:
    - set_prompt(): Wraps prompt text in OpenTelemetry message structure
    - set_completion(): Wraps completion text in OpenTelemetry message structure

    All other GenAI methods (set_tokens, set_model, etc.) are available via the base
    SpanHandle class and can be used with any span type.

    Follows OpenTelemetry GenAI semantic conventions.
    See: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/
    """

    def set_prompt(self, prompt: str | Prompt) -> None:
        """
        Set the prompt/input text.
        Note: This uses gen_ai.input.messages for structured messages.
        For simple string prompts, wraps in a message structure.
        """
        if not isinstance(prompt, str):
            # Assume it's a Prompt object
            super().set_prompt(prompt)
            prompt_text = prompt.text
        else:
            prompt_text = prompt

        super().set_input(prompt_text)
        if trace_content_enabled():
            # Store as structured message format per OpenTelemetry spec
            messages = [{"role": "user", "parts": [{"type": "text", "content": prompt_text}]}]
            self.set_attribute(semconv.GenAI.INPUT_MESSAGES, json.dumps(messages))

    def set_completion(self, completion: str) -> None:
        """
        Set the completion/output text.
        Note: This uses gen_ai.output.messages for structured messages.
        For simple string completions, wraps in a message structure.
        """
        super().set_output(completion)
        if trace_content_enabled():
            # Store as structured message format per OpenTelemetry spec
            messages = [
                {
                    "role": "assistant",
                    "parts": [{"type": "text", "content": completion}],
                    "finish_reason": "stop",
                }
            ]
            self.set_attribute(semconv.GenAI.OUTPUT_MESSAGES, json.dumps(messages))


class RetrievalSpanHandle(SpanHandle):
    """
    Span handle for vector DB/retrieval operations.

    Uses Basalt-specific semantic conventions for retrieval operations.
    """

    def set_query(self, query: str) -> None:
        """Set the query text for the retrieval operation."""
        super().set_input(query)
        self.set_attribute(semconv.BasaltRetrieval.QUERY, query)

    def set_results_count(self, count: int) -> None:
        """Set the number of results returned."""
        self.set_attribute(semconv.BasaltRetrieval.RESULTS_COUNT, count)

    def set_top_k(self, top_k: float) -> None:
        """Set the top-K parameter for retrieval."""
        value = int(top_k) if isinstance(top_k, float) else top_k
        self.set_attribute(semconv.BasaltRetrieval.TOP_K, value)


class FunctionSpanHandle(SpanHandle):
    """
    Span handle for compute/function execution spans.
    """

    def set_function_name(self, function_name: str) -> None:
        """Set the logical function name being executed."""
        self.set_attribute(semconv.BasaltFunction.NAME, function_name)

    def set_stage(self, stage: str) -> None:
        """Set the stage or phase associated with the execution."""
        self.set_attribute(semconv.BasaltFunction.STAGE, stage)

    def add_metric(self, key: str, value: str | int | float | bool) -> None:
        """Attach custom metric data to the function execution."""
        self.set_attribute(f"{semconv.BasaltFunction.METRIC_PREFIX}.{key}", value)


class ToolSpanHandle(SpanHandle):
    """
    Span handle for tool invocation spans.

    Uses Basalt-specific semantic conventions for tool operations.
    """

    def set_tool_name(self, name: str) -> None:
        """Set the name of the tool being invoked."""
        self.set_attribute(semconv.BasaltTool.NAME, name)

    def set_input(self, payload: JSONValue) -> None:
        """Set the input payload for the tool."""
        super().set_input(payload)
        if trace_content_enabled():
            value = _serialize_attribute(payload)
            if value is not None:
                self.set_attribute(semconv.BasaltTool.INPUT, value)

    def set_output(self, payload: JSONValue) -> None:
        """Set the output payload from the tool."""
        super().set_output(payload)
        if trace_content_enabled():
            value = _serialize_attribute(payload)
            if value is not None:
                self.set_attribute(semconv.BasaltTool.OUTPUT, value)


class EventSpanHandle(SpanHandle):
    """
    Span handle for custom application events.

    Uses Basalt-specific semantic conventions for event operations.
    """

    def set_event_type(self, event_type: str) -> None:
        """Set the type of custom event."""
        self.set_attribute(semconv.BasaltEvent.TYPE, event_type)

    def set_payload(self, payload: JSONValue) -> None:
        """Set the event payload."""
        super().set_input(payload)
        if trace_content_enabled():
            value = _serialize_attribute(payload)
            if value is not None:
                self.set_attribute(semconv.BasaltEvent.PAYLOAD, value)


@contextmanager
def _with_span_handle(
    name: str,
    attributes: dict[str, Any] | None,
    tracer_name: str,
    handle_cls: type[SpanHandle],
    span_type: str | None = None,
    *,
    input_payload: JSONValue | None = None,
    output_payload: JSONValue | None = None,
    variables: Mapping[str, Any] | None = None,
    evaluators: Sequence[Any] | None = None,
    user: TraceIdentity | Mapping[str, Any] | None = None,
    organization: TraceIdentity | Mapping[str, Any] | None = None,
    feature_slug: str | None = None,
    metadata: Mapping[str, Any] | None = None,
    evaluate_config: EvaluationConfig | None = None,
    experiment: str | dict[str, Any] | None = None,
) -> Generator[SpanHandle, None, None]:
    tracer = get_tracer(tracer_name)
    defaults = _current_trace_defaults()

    parent_span = trace.get_current_span()
    if parent_span and (not parent_span.get_span_context().is_valid or not parent_span.is_recording()):
        parent_span = None

    # Prepare context tokens for user/org propagation
    tokens = []
    if user is not None:
        from .trace_context import _coerce_identity

        user_identity = _coerce_identity(user)
        if user_identity:
            tokens.append(attach(set_value(USER_CONTEXT_KEY, user_identity)))

    if organization is not None:
        from .trace_context import _coerce_identity

        org_identity = _coerce_identity(organization)
        if org_identity:
            tokens.append(attach(set_value(ORGANIZATION_CONTEXT_KEY, org_identity)))

    if feature_slug is not None:
        from .trace_context import FEATURE_SLUG_CONTEXT_KEY

        tokens.append(attach(set_value(FEATURE_SLUG_CONTEXT_KEY, feature_slug)))

    if experiment is not None:
        from .trace_context import EXPERIMENT_CONTEXT_KEY

        tokens.append(attach(set_value(EXPERIMENT_CONTEXT_KEY, experiment)))

    # Determine if this should be treated as a Basalt root span
    # A span is a Basalt root if:
    # 1. There's no parent span at all (true root), OR
    # 2. This is a start_observe (span_type="basalt_trace") AND the parent is NOT a Basalt span
    #    (e.g., parent is from FastAPI, httpx, or other instrumentation)
    root_span_token = None
    root_parent_span_token = None

    # Check if we're already inside a basalt trace
    in_basalt_trace = otel_context.get_value(ROOT_SPAN_CONTEXT_KEY) is not None

    # Determine root status
    if parent_span is None:
        # No parent at all - this is a true root
        is_root = True
    elif span_type == "basalt_trace" and not in_basalt_trace:
        # Parent exists but it's NOT a Basalt span (e.g., FastAPI HTTP span)
        # AND this is a start_observe call (basalt_trace type)
        # -> Treat as Basalt root (allows start_observe to work inside FastAPI handlers)
        is_root = True

        # Retroactively apply basalt trace metadata to the parent span
        # This ensures the entire trace tree has consistent basalt attributes
        if parent_span is not None and parent_span.is_recording():
            parent_span.set_attribute(semconv.BasaltSpan.IN_TRACE, True)

            parent_span.set_attribute(semconv.BasaltSpan.KIND, "span")

            # Also propagate feature_slug to parent if available
            if feature_slug is not None:
                parent_span.set_attribute(semconv.BasaltSpan.FEATURE_SLUG, feature_slug)

            # Also propagate experiment to parent if available
            if experiment is not None:
                from .processors import _apply_experiment_attrs

                _apply_experiment_attrs(parent_span, experiment)
    else:
        # Parent exists and either:
        # - We're already in a Basalt trace, OR
        # - This is not a start_observe call
        # -> Treat as nested span
        is_root = False

    # Make trace-level sampling decision
    should_evaluate_token = None
    if is_root:
        # Root span: make new sampling decision
        # If experiment is attached, ALWAYS evaluate (should_evaluate=True)
        if experiment is not None:
            should_evaluate = True
        else:
            # Get sample_rate from evaluate_config if provided, otherwise use global default
            if evaluate_config is not None:
                effective_sample_rate = evaluate_config.sample_rate
            else:
                effective_sample_rate = defaults.sample_rate
            should_evaluate = random.random() < effective_sample_rate
        should_evaluate_token = attach(set_value(SHOULD_EVALUATE_CONTEXT_KEY, should_evaluate))
    else:
        # Check if should_evaluate already exists in context
        existing_should_evaluate = otel_context.get_value(SHOULD_EVALUATE_CONTEXT_KEY)
        if existing_should_evaluate is None:
            # Orphan span without root - make its own decision
            # If experiment is attached, ALWAYS evaluate
            if experiment is not None:
                should_evaluate = True
            else:
                if evaluate_config is not None:
                    effective_sample_rate = evaluate_config.sample_rate
                else:
                    effective_sample_rate = defaults.sample_rate
                should_evaluate = random.random() < effective_sample_rate
            should_evaluate_token = attach(set_value(SHOULD_EVALUATE_CONTEXT_KEY, should_evaluate))

    try:
        with tracer.start_as_current_span(name) as span:
            # Store root span in context for retrieval from nested spans
            if is_root:
                root_span_token = attach(set_value(ROOT_SPAN_CONTEXT_KEY, span))
                # Also store the parent span (if any) for trace-level helpers
                if parent_span is not None:
                    root_parent_span_token = attach(set_value(ROOT_PARENT_SPAN_CONTEXT_KEY, parent_span))
                # Set basalt.root attribute
                span.set_attribute("basalt.root", True)
            elif in_basalt_trace:
                # Child span inside a basalt trace
                span.set_attribute("basalt.trace", True)

            # Mark all basalt spans with basalt.in_trace
            span.set_attribute(semconv.BasaltSpan.IN_TRACE, True)

            # Inject prompt context if available
            try:
                from basalt.prompts.models import _current_prompt_context

                prompt_ctx = _current_prompt_context.get()
                if prompt_ctx:
                    from .utils import apply_prompt_context_attributes

                    apply_prompt_context_attributes(span, prompt_ctx)
            except ImportError:
                # Prompts module not available, skip injection
                pass

            _attach_attributes(span, attributes)

            # Apply metadata if provided
            if metadata:
                from .utils import apply_span_metadata

                apply_span_metadata(span, metadata)

            if span_type:
                span.set_attribute(SPAN_TYPE_ATTRIBUTE, span_type)

            # Apply user/org from context (either explicit or inherited from parent)
            apply_user_from_context(span, user)
            apply_organization_from_context(span, organization)

            if is_root and handle_cls == SpanHandle:
                actual_handle_cls = StartSpanHandle
            else:
                actual_handle_cls = handle_cls

            handle = actual_handle_cls(span, parent_span, defaults)
            if input_payload is not None:
                handle.set_input(input_payload)
            if variables:
                handle.set_io(variables=variables)
            if evaluators:
                handle.add_evaluators(*evaluators)
            yield handle
            if output_payload is not None:
                handle.set_output(output_payload)

    finally:
        # Detach should_evaluate token if it was set
        if should_evaluate_token is not None:
            detach(should_evaluate_token)

        # Detach root span tokens if they were set
        if root_parent_span_token is not None:
            detach(root_parent_span_token)
        if root_span_token is not None:
            detach(root_span_token)

        # Detach context tokens in reverse order
        for token in reversed(tokens):
            detach(token)


@asynccontextmanager
async def _async_with_span_handle(
    name: str,
    attributes: dict[str, Any] | None,
    tracer_name: str,
    handle_cls: type[SpanHandle],
    span_type: str | None = None,
    *,
    input_payload: JSONValue | None = None,
    output_payload: JSONValue | None = None,
    variables: Mapping[str, Any] | None = None,
    evaluators: Sequence[Any] | None = None,
    user: TraceIdentity | Mapping[str, Any] | None = None,
    organization: TraceIdentity | Mapping[str, Any] | None = None,
    feature_slug: str | None = None,
    metadata: Mapping[str, Any] | None = None,
    evaluate_config: EvaluationConfig | None = None,
    experiment: str | dict[str, Any] | Experiment | None = None,
) -> AsyncGenerator[SpanHandle, None]:
    """Async version of _with_span_handle.

    Note: OpenTelemetry's tracer.start_as_current_span() is synchronous,
    so this async context manager still calls sync OTel APIs internally.
    The async support is primarily for use with async with statements.
    """
    # Coerce Experiment to str or dict if needed
    experiment_arg: str | dict[str, Any] | None = None
    if experiment is not None:
        if isinstance(experiment, (str, dict)):
            experiment_arg = experiment
        else:
            experiment_arg = getattr(experiment, "id", str(experiment))

    with _with_span_handle(
        name=name,
        attributes=attributes,
        tracer_name=tracer_name,
        handle_cls=handle_cls,
        span_type=span_type,
        input_payload=input_payload,
        output_payload=output_payload,
        variables=variables,
        evaluators=evaluators,
        user=user,
        organization=organization,
        feature_slug=feature_slug,
        metadata=metadata,
        evaluate_config=evaluate_config,
        experiment=experiment_arg,
    ) as handle:
        yield handle


def _set_trace_user(user_id: str, name: str | None = None) -> None:
    """Set the user identity for the current trace context."""
    identity = TraceIdentity(id=user_id, name=name)
    attach(set_value(USER_CONTEXT_KEY, identity))


def _set_trace_organization(organization_id: str, name: str | None = None) -> None:
    """Set the organization identity for the current trace context."""
    identity = TraceIdentity(id=organization_id, name=name)
    attach(set_value(ORGANIZATION_CONTEXT_KEY, identity))
