rialo-telemetry 0.3.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! OpenTelemetry-specific functionality for distributed tracing.

use std::collections::HashMap;

use anyhow::Result;
use opentelemetry::{
    propagation::{Extractor, Injector},
    trace::TracerProvider,
};
use opentelemetry_otlp::{SpanExporter, WithExportConfig, WithHttpConfig};
use opentelemetry_sdk::{trace as sdktrace, trace::Sampler};
use tracing_opentelemetry::OpenTelemetrySpanExt;

use crate::{
    parse_env::{parse_bool_env, parse_key_value_env, parse_string_list_env},
    TelemetryHandle,
};

/// Default OTLP endpoint for general export (HTTP)
pub const DEFAULT_OTLP_ENDPOINT: &str = "http://localhost:4318";

/// Default OTLP endpoint for gRPC export
pub const DEFAULT_OTLP_GRPC_ENDPOINT: &str = "http://localhost:4317";

/// Result of OpenTelemetry initialization containing both the handle and tracer
pub struct OtelResult {
    pub handle: TelemetryHandle,
    pub tracer: Option<sdktrace::SdkTracer>,
}

/// Protocol for OTLP export
#[derive(Debug, Clone, PartialEq, Default)]
pub enum Protocol {
    Grpc,
    #[default]
    HttpProtobuf,
    HttpJson,
}

/// Sampling frequency of spans.
#[derive(Debug, Clone, PartialEq)]
pub enum Sampling {
    AlwaysOn,
    AlwaysOff,
    Ratio(f64),
}

impl Default for Sampling {
    fn default() -> Self {
        Sampling::Ratio(0.05)
    }
}

/// OTLP-specific configuration
#[derive(Debug, Clone)]
pub struct OtlpConfig {
    /// Service name (OTEL_SERVICE_NAME, default: "rialo")
    pub service_name: String,
    /// Service version (OTEL_SERVICE_VERSION, default: "unknown")
    pub service_version: String,
    /// Log level (OTEL_LOG_LEVEL, default: "info")
    pub log_level: String,
    /// Resource attributes (OTEL_RESOURCE_ATTRIBUTES)
    pub resource_attributes: HashMap<String, String>,
    /// Whether to enable console output (not from env, default: true)
    pub enable_console: bool,
    /// General OTLP exporter endpoint (OTEL_EXPORTER_OTLP_ENDPOINT, default: DEFAULT_OTLP_ENDPOINT)
    pub exporter_endpoint: Option<String>,
    /// Whether general endpoint is insecure (OTEL_EXPORTER_OTLP_INSECURE, default: false)
    pub exporter_endpoint_insecure: bool,
    /// Traces-specific exporter endpoint (OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, default: DEFAULT_TRACES_ENDPOINT)
    pub traces_exporter_endpoint: Option<String>,
    /// Whether traces endpoint is insecure (OTEL_EXPORTER_OTLP_TRACES_INSECURE, default: false)
    pub traces_exporter_endpoint_insecure: bool,
    /// Whether traces are enabled (OTEL_TRACES_ENABLED, default: true)
    pub traces_enabled: bool,
    /// Metrics exporter endpoint (OTEL_EXPORTER_OTLP_METRICS_ENDPOINT, default: DEFAULT_METRICS_ENDPOINT)
    ///
    /// **NOTE: Future Work** - This configuration is parsed for compatibility but not currently
    /// used for actual metrics export. OpenTelemetry metrics implementation is planned for
    /// a future release. Use the `prometheus` feature for metrics collection instead.
    pub metrics_exporter_endpoint: Option<String>,
    /// Whether metrics endpoint is insecure (OTEL_EXPORTER_OTLP_METRICS_INSECURE, default: false)
    ///
    /// **NOTE: Future Work** - Parsed but not implemented. See metrics_exporter_endpoint.
    pub metrics_exporter_endpoint_insecure: bool,
    /// Whether metrics are enabled (OTEL_METRICS_ENABLED, default: true)
    ///
    /// **NOTE: Future Work** - Parsed but not implemented. See metrics_exporter_endpoint.
    pub metrics_enabled: bool,
    /// Metrics reporting period (OTEL_EXPORTER_OTLP_METRICS_PERIOD, default: "30s")
    ///
    /// **NOTE: Future Work** - Parsed but not implemented. See metrics_exporter_endpoint.
    pub metrics_reporting_period: String,
    /// Propagators (OTEL_PROPAGATORS, default: ["tracecontext", "baggage"])
    pub propagators: Vec<String>,
    /// General exporter protocol (OTEL_EXPORTER_OTLP_PROTOCOL, default: http/protobuf)
    pub exporter_protocol: Protocol,
    /// Traces exporter protocol (OTEL_EXPORTER_OTLP_TRACES_PROTOCOL)
    pub traces_exporter_protocol: Option<Protocol>,
    /// Metrics exporter protocol (OTEL_EXPORTER_OTLP_METRICS_PROTOCOL)
    ///
    /// **NOTE: Future Work** - Parsed but not implemented. See metrics_exporter_endpoint.
    pub metrics_exporter_protocol: Option<Protocol>,
    /// General headers (OTEL_EXPORTER_OTLP_HEADERS)
    pub headers: HashMap<String, String>,
    /// Traces headers (OTEL_EXPORTER_OTLP_TRACES_HEADERS)
    pub traces_headers: HashMap<String, String>,
    /// Metrics headers (OTEL_EXPORTER_OTLP_METRICS_HEADERS)
    ///
    /// **NOTE: Future Work** - Parsed but not implemented. See metrics_exporter_endpoint.
    pub metrics_headers: HashMap<String, String>,
    /// Controls sampling of traces. (OTEL_TRACES_SAMPLER: "always_on", "always_off", "traceidratio" with OTEL_TRACES_SAMPLER_ARG="0.1")
    pub sampling: Sampling,
}

impl Default for OtlpConfig {
    fn default() -> Self {
        // Get the main protocol first
        let main_protocol = parse_protocol_env("OTEL_EXPORTER_OTLP_PROTOCOL", Protocol::default());

        // Check if traces-specific protocol is set (takes precedence for determining default endpoint)
        let traces_protocol = std::env::var("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL")
            .ok()
            .map(|_| {
                parse_protocol_env("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", main_protocol.clone())
            })
            .unwrap_or_else(|| main_protocol.clone());

        // Choose default endpoint based on the effective traces protocol
        let default_endpoint = match traces_protocol {
            Protocol::Grpc => DEFAULT_OTLP_GRPC_ENDPOINT,
            _ => DEFAULT_OTLP_ENDPOINT,
        };

        // Get the general endpoint first (from env or default based on protocol)
        let general_endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT")
            .ok()
            .or_else(|| Some(default_endpoint.to_string()));

        Self {
            service_name: std::env::var("OTEL_SERVICE_NAME")
                .unwrap_or_else(|_| "rialo".to_string()),
            service_version: std::env::var("OTEL_SERVICE_VERSION")
                .unwrap_or_else(|_| "unknown".to_string()),
            log_level: std::env::var("OTEL_LOG_LEVEL").unwrap_or_else(|_| "info".to_string()),
            resource_attributes: parse_key_value_env("OTEL_RESOURCE_ATTRIBUTES"),
            enable_console: true,
            exporter_endpoint: general_endpoint.clone(),
            exporter_endpoint_insecure: parse_bool_env("OTEL_EXPORTER_OTLP_INSECURE", false),
            traces_exporter_endpoint: std::env::var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT").ok(),
            traces_exporter_endpoint_insecure: parse_bool_env(
                "OTEL_EXPORTER_OTLP_TRACES_INSECURE",
                false,
            ),
            traces_enabled: parse_bool_env("OTEL_TRACES_ENABLED", true),
            metrics_exporter_endpoint: std::env::var("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT").ok(),
            metrics_exporter_endpoint_insecure: parse_bool_env(
                "OTEL_EXPORTER_OTLP_METRICS_INSECURE",
                false,
            ),
            metrics_enabled: parse_bool_env("OTEL_METRICS_ENABLED", true),
            metrics_reporting_period: std::env::var("OTEL_EXPORTER_OTLP_METRICS_PERIOD")
                .unwrap_or_else(|_| "30s".to_string()),
            propagators: parse_string_list_env(
                "OTEL_PROPAGATORS",
                vec!["tracecontext".to_string(), "baggage".to_string()],
            ),
            exporter_protocol: main_protocol.clone(),
            traces_exporter_protocol: std::env::var("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL")
                .ok()
                .map(|_| {
                    parse_protocol_env("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", main_protocol.clone())
                }),
            metrics_exporter_protocol: std::env::var("OTEL_EXPORTER_OTLP_METRICS_PROTOCOL")
                .ok()
                .map(|_| parse_protocol_env("OTEL_EXPORTER_OTLP_METRICS_PROTOCOL", main_protocol)),
            headers: parse_key_value_env("OTEL_EXPORTER_OTLP_HEADERS"),
            traces_headers: parse_key_value_env("OTEL_EXPORTER_OTLP_TRACES_HEADERS"),
            metrics_headers: parse_key_value_env("OTEL_EXPORTER_OTLP_METRICS_HEADERS"),
            sampling: parse_sampler_env(
                "OTEL_TRACES_SAMPLER",
                "OTEL_TRACES_SAMPLER_ARG",
                Sampling::default(),
            ),
        }
    }
}

impl OtlpConfig {
    /// Create a new OtlpConfig with default values from environment variables.
    pub fn new() -> Self {
        Self::default()
    }

    /// Validate URL format
    fn validate_url(url: &str) -> anyhow::Result<()> {
        if url.is_empty() {
            return Err(anyhow::anyhow!("URL cannot be empty. Provide a valid endpoint URL or omit the configuration to disable tracing."));
        }

        url.parse::<url::Url>()
            .map_err(|_| anyhow::anyhow!("Invalid URL format: '{}'. Please provide a valid URL (e.g., 'http://localhost:4318' or 'https://api.honeycomb.io/v1/traces')", url))?;

        Ok(())
    }

    /// Validate service name format
    fn validate_service_name(name: &str) -> anyhow::Result<()> {
        if name.is_empty() {
            return Err(anyhow::anyhow!("Service name cannot be empty"));
        }

        if name.len() > 255 {
            return Err(anyhow::anyhow!("Service name cannot exceed 255 characters"));
        }

        // Basic character validation - alphanumeric, hyphens, underscores, dots
        if !name
            .chars()
            .all(|c| c.is_alphanumeric() || c == '-' || c == '_' || c == '.')
        {
            return Err(anyhow::anyhow!("Service name can only contain alphanumeric characters, hyphens, underscores, and dots"));
        }

        Ok(())
    }

    /// Validate header value format
    fn validate_header_value(value: &str) -> anyhow::Result<()> {
        // HTTP header values must be ASCII and cannot contain control characters
        if !value.is_ascii() {
            return Err(anyhow::anyhow!("Header values must be ASCII"));
        }

        if value.chars().any(|c| c.is_control() && c != '\t') {
            return Err(anyhow::anyhow!(
                "Header values cannot contain control characters (except tab)"
            ));
        }

        Ok(())
    }

    /// Validate the configuration and return errors if any issues are found
    pub fn validate(&self) -> anyhow::Result<()> {
        Self::validate_service_name(&self.service_name)?;

        if let Some(ref endpoint) = self.exporter_endpoint {
            Self::validate_url(endpoint)?;
        }

        if let Some(ref endpoint) = self.traces_exporter_endpoint {
            Self::validate_url(endpoint)?;
        }

        // Validate headers
        for value in self.headers.values() {
            Self::validate_header_value(value)?;
        }

        for value in self.traces_headers.values() {
            Self::validate_header_value(value)?;
        }

        Ok(())
    }

    /// Set the service name.
    pub fn with_service_name(mut self, name: impl Into<String>) -> Self {
        self.service_name = name.into();
        self
    }

    /// Set the service version.
    pub fn with_service_version(mut self, version: impl Into<String>) -> Self {
        self.service_version = version.into();
        self
    }

    /// Set the log level.
    pub fn with_log_level(mut self, level: impl Into<String>) -> Self {
        self.log_level = level.into();
        self
    }

    /// Enable or disable console output.
    pub fn with_console_enabled(mut self, enabled: bool) -> Self {
        self.enable_console = enabled;
        self
    }

    /// Set whether traces are enabled.
    pub fn with_traces_enabled(mut self, enabled: bool) -> Self {
        self.traces_enabled = enabled;
        self
    }

    /// Set the general exporter endpoint.
    pub fn with_exporter_endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.exporter_endpoint = Some(endpoint.into());
        self
    }

    /// Set the traces exporter endpoint.
    pub fn with_traces_endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.traces_exporter_endpoint = Some(endpoint.into());
        self
    }

    pub fn with_exporter_protocol(mut self, protocol: Protocol) -> Self {
        self.exporter_protocol = protocol;
        self
    }

    pub fn with_traces_sampling(mut self, sampling: Sampling) -> Self {
        self.sampling = sampling;
        self
    }

    pub fn add_resource_attribute(mut self, key: String, value: String) -> Self {
        self.resource_attributes.insert(key, value);
        self
    }

    /// Get the effective traces endpoint (traces-specific takes precedence over general)
    pub fn effective_traces_endpoint(&self) -> Option<&String> {
        self.traces_exporter_endpoint
            .as_ref()
            .or(self.exporter_endpoint.as_ref())
    }

    /// Get the effective traces protocol (traces-specific takes precedence over general)
    pub fn effective_traces_protocol(&self) -> &Protocol {
        self.traces_exporter_protocol
            .as_ref()
            .unwrap_or(&self.exporter_protocol)
    }

    /// Get the effective traces headers (general headers + traces-specific headers)
    /// This method returns an owned HashMap for compatibility but consider using
    /// `iter_effective_traces_headers` for read-heavy scenarios to avoid cloning.
    pub fn effective_traces_headers(&self) -> HashMap<String, String> {
        // Optimize for the common case where one of the maps is empty
        if self.headers.is_empty() {
            return self.traces_headers.clone();
        }
        if self.traces_headers.is_empty() {
            return self.headers.clone();
        }

        // Both maps have content, need to merge
        let mut headers = self.headers.clone();
        headers.extend(self.traces_headers.clone());
        headers
    }

    /// Iterate over effective traces headers without creating a new HashMap.
    /// This is more efficient for read-only operations.
    pub fn iter_effective_traces_headers(&self) -> impl Iterator<Item = (&String, &String)> {
        self.headers.iter().chain(self.traces_headers.iter())
    }

    /// Check if a header exists in the effective traces headers without cloning
    pub fn has_effective_traces_header(&self, key: &str) -> bool {
        self.headers.contains_key(key) || self.traces_headers.contains_key(key)
    }

    /// Get a header value from effective traces headers without cloning the entire map
    pub fn get_effective_traces_header(&self, key: &str) -> Option<&String> {
        self.traces_headers
            .get(key)
            .or_else(|| self.headers.get(key))
    }
}

/// Initialize OpenTelemetry distributed tracing.
///
/// This function sets up the OpenTelemetry tracer with OTLP HTTP exporter
/// and returns both a handle to the tracer provider and the tracer itself.
pub async fn init_otel(otlp_config: &OtlpConfig) -> Result<OtelResult> {
    // Validate configuration first
    otlp_config
        .validate()
        .map_err(|e| anyhow::anyhow!("OTLP configuration validation failed: {}", e))?;

    // Early return if traces are disabled
    if !otlp_config.traces_enabled {
        return Ok(OtelResult {
            handle: TelemetryHandle::empty(),
            tracer: None,
        });
    }

    // Create resource with attributes directly (avoid environment variable mutation)
    let mut resource_builder = opentelemetry_sdk::resource::Resource::builder()
        .with_service_name(otlp_config.service_name.clone())
        .with_attribute(opentelemetry::KeyValue::new(
            "service.version",
            otlp_config.service_version.clone(),
        ));

    // Add custom resource attributes
    for (key, value) in &otlp_config.resource_attributes {
        resource_builder = resource_builder
            .with_attribute(opentelemetry::KeyValue::new(key.clone(), value.clone()));
    }

    let otel_resource = resource_builder.build();

    // Check if we have a valid endpoint for OpenTelemetry
    let endpoint = match otlp_config.effective_traces_endpoint() {
        Some(endpoint) if !endpoint.is_empty() => endpoint,
        Some(_) => {
            return Err(anyhow::anyhow!(
                "Empty endpoint URL provided. Either provide a valid endpoint URL or disable tracing by setting traces_enabled=false"
            ));
        }
        None => {
            return Ok(OtelResult {
                handle: TelemetryHandle::empty(),
                tracer: None,
            });
        }
    };

    let protocol = otlp_config
        .traces_exporter_protocol
        .clone()
        .unwrap_or_else(|| otlp_config.exporter_protocol.clone());

    if otlp_config.traces_exporter_protocol.is_some()
        && otlp_config.traces_exporter_protocol != Some(otlp_config.exporter_protocol.clone())
    {
        println!(
            "Using OTLP traces protocol override: {:?}",
            otlp_config.traces_exporter_protocol
        );
    }

    if matches!(protocol, Protocol::HttpProtobuf | Protocol::HttpJson)
        && !endpoint.ends_with("/v1/traces")
    {
        println!(
            "OTLP HTTP endpoint does not end with /v1/traces: {}",
            endpoint
        );
    }

    // Set up OpenTelemetry exporter
    let exporter = match protocol {
        Protocol::Grpc => SpanExporter::builder()
            .with_tonic()
            .with_endpoint(endpoint)
            .build()?,
        Protocol::HttpProtobuf | Protocol::HttpJson => {
            let mut exporter_builder = SpanExporter::builder().with_http().with_endpoint(endpoint);
            // Add effective headers
            let headers = otlp_config.effective_traces_headers();
            if !headers.is_empty() {
                exporter_builder = exporter_builder.with_headers(headers);
            }
            exporter_builder.build()?
        }
    };

    let sampler = match otlp_config.sampling {
        Sampling::AlwaysOn => Sampler::AlwaysOn,
        Sampling::AlwaysOff => Sampler::AlwaysOff,
        Sampling::Ratio(x) => Sampler::TraceIdRatioBased(x),
    };

    // Create tracer provider with resource
    let tracer_provider = sdktrace::SdkTracerProvider::builder()
        .with_batch_exporter(exporter)
        .with_resource(otel_resource)
        .with_sampler(sampler)
        .build();

    // Create tracer from provider
    let tracer = tracer_provider.tracer(otlp_config.service_name.clone());

    // Set as global tracer provider
    opentelemetry::global::set_tracer_provider(tracer_provider.clone());

    // Set up text map propagator for distributed tracing
    let propagator = opentelemetry_sdk::propagation::TraceContextPropagator::new();
    opentelemetry::global::set_text_map_propagator(propagator);

    Ok(OtelResult {
        handle: TelemetryHandle::new(tracer_provider),
        tracer: Some(tracer),
    })
}

/// Simple header injector for OpenTelemetry trace context propagation
struct HeaderInjector<'a>(&'a mut HashMap<String, String>);

impl<'a> Injector for HeaderInjector<'a> {
    fn set(&mut self, key: &str, value: String) {
        self.0.insert(key.to_string(), value);
    }
}

/// Inject OpenTelemetry trace context into HTTP headers for distributed tracing
///
/// This function extracts trace context from the current span and returns a HashMap
/// of headers that should be added to HTTP requests for distributed tracing.
///
/// # Returns
///
/// A HashMap containing trace context headers that should be injected into HTTP requests
///
/// # Examples
///
/// ```rust,ignore
/// use rialo_telemetry::inject_trace_headers;
///
/// let headers = inject_trace_headers();
/// // Add headers to your HTTP request
/// for (key, value) in headers {
///     request = request.header(key, value);
/// }
/// ```
pub fn inject_trace_headers() -> HashMap<String, String> {
    // Get current span context and inject it into HTTP headers
    let context = tracing::Span::current().context();
    let mut headers = HashMap::new();

    opentelemetry::global::get_text_map_propagator(|propagator| {
        propagator.inject_context(&context, &mut HeaderInjector(&mut headers));
    });

    headers
}

/// Convenience function for applying trace headers to a reqwest RequestBuilder
///
/// This function takes the headers from `inject_trace_headers()` and applies them
/// to a reqwest RequestBuilder for distributed tracing across HTTP requests.
///
/// # Arguments
///
/// * `request_builder` - The reqwest RequestBuilder to inject headers into
/// * `headers` - The headers from `inject_trace_headers()`
///
/// # Returns
///
/// The RequestBuilder with trace context headers added
///
/// # Examples
///
/// ```rust,ignore
/// use rialo_telemetry::opentelemetry::{inject_trace_headers, apply_trace_headers_to_reqwest};
///
/// let headers = inject_trace_headers();
/// let request = client.post("https://api.example.com");
/// let request = apply_trace_headers_to_reqwest(request, headers);
/// let response = request.send().await?;
/// ```
pub fn apply_trace_headers_to_reqwest(
    mut request_builder: reqwest::RequestBuilder,
    headers: HashMap<String, String>,
) -> reqwest::RequestBuilder {
    for (key, value) in headers {
        if let (Ok(header_name), Ok(header_value)) = (
            reqwest::header::HeaderName::from_bytes(key.as_bytes()),
            reqwest::header::HeaderValue::from_str(&value),
        ) {
            tracing::trace!("Adding trace header: {} = {}", key, value);
            request_builder = request_builder.header(header_name, header_value);
        }
    }
    request_builder
}

/// Extract OpenTelemetry trace context from axum HeaderMap and set it on the current span
///
/// This function extracts trace context from axum HTTP headers and sets it as the parent
/// context for the current tracing span, enabling distributed tracing across service boundaries.
///
/// # Arguments
///
/// * `headers` - The axum HeaderMap containing trace context information
///
/// # Examples
///
/// ```rust,ignore
/// use rialo_telemetry::opentelemetry::extract_and_set_trace_context_axum;
///
/// #[tracing::instrument]
/// async fn handler(headers: HeaderMap) -> impl IntoResponse {
///     extract_and_set_trace_context_axum(&headers);
///     // Your handler logic here
/// }
/// ```
pub fn extract_and_set_trace_context_axum(headers: &axum::http::HeaderMap) {
    /// Local header extractor for axum HeaderMap
    struct AxumHeaderExtractor<'a>(&'a axum::http::HeaderMap);

    impl<'a> Extractor for AxumHeaderExtractor<'a> {
        fn get(&self, key: &str) -> Option<&str> {
            self.0.get(key).and_then(|value| value.to_str().ok())
        }

        fn keys(&self) -> Vec<&str> {
            self.0.keys().map(|name| name.as_str()).collect()
        }
    }

    tracing::debug!("Extracting trace context from {} headers", headers.len());

    let extractor = AxumHeaderExtractor(headers);
    let context =
        opentelemetry::global::get_text_map_propagator(|propagator| propagator.extract(&extractor));

    let span = tracing::Span::current();
    tracing::debug!("Setting parent context on span: {:?}", span.id());
    let _ = span.set_parent(context);
}

/// Simple environment variable injector for OpenTelemetry trace context propagation
struct EnvInjector<'a>(&'a mut HashMap<String, String>);

impl<'a> Injector for EnvInjector<'a> {
    fn set(&mut self, key: &str, value: String) {
        self.0.insert(key.to_string(), value);
    }
}

/// Simple environment variable extractor for OpenTelemetry trace context propagation
struct EnvExtractor<'a>(&'a HashMap<String, String>);

impl<'a> Extractor for EnvExtractor<'a> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).map(|s| s.as_str())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(|k| k.as_str()).collect()
    }
}

/// Inject OpenTelemetry trace context into environment variables for subprocess propagation
///
/// This function extracts trace context from the current span and returns a HashMap
/// of environment variables that should be passed to child processes/commands for
/// distributed tracing across process boundaries.
///
/// # Returns
///
/// A HashMap containing environment variables with trace context that should be passed to subprocesses
///
/// # Examples
///
/// ```rust,ignore
/// use rialo_telemetry::inject_trace_env;
/// use std::process::Command;
///
/// // In a traced context
/// let span = tracing::info_span!("subprocess_execution");
/// let _guard = span.enter();
///
/// let trace_env = inject_trace_env();
/// let mut cmd = Command::new("my-subprocess");
///
/// // Add trace context environment variables
/// for (key, value) in trace_env {
///     cmd.env(key, value);
/// }
///
/// let output = cmd.output().expect("Failed to execute command");
/// ```
pub fn inject_trace_env() -> HashMap<String, String> {
    // Get current span context and inject it into environment variables
    let context = tracing::Span::current().context();
    let mut env_vars = HashMap::new();

    opentelemetry::global::get_text_map_propagator(|propagator| {
        propagator.inject_context(&context, &mut EnvInjector(&mut env_vars));
    });

    env_vars
}

/// Extract OpenTelemetry trace context from environment variables and set up distributed tracing
///
/// This function extracts trace context from the current process's environment variables
/// and sets up the tracing context to inherit from a parent process, enabling distributed
/// tracing across process boundaries.
///
/// If no active span exists when trace context is found, a root span will be created
/// to represent the entire process lifetime. This span will be entered and kept active
/// for the duration of the process, and the OpenTelemetry context will be set globally
/// to ensure all future spans inherit the correct trace context.
///
/// # Examples
///
/// ```rust,ignore
/// use rialo_telemetry::extract_and_set_trace_context_env;
///
/// fn main() {
///     // Initialize telemetry first
///     let _handle = rialo_telemetry::init_telemetry(config).await;
///     
///     // This will automatically extract trace context and set up distributed tracing
///     // No need to manage return values - it's all handled automatically
///     
///     // Any spans created after this will be part of the parent process's distributed trace
///     tracing::info!("Subprocess started with inherited trace context");
/// }
/// ```
pub fn extract_and_set_trace_context_env() {
    // Collect relevant environment variables for trace context
    let env_vars: HashMap<String, String> = std::env::vars()
        .filter(|(key, _)| {
            // Only collect OpenTelemetry trace context related environment variables
            key.starts_with("traceparent")
                || key.starts_with("tracestate")
                || key.starts_with("baggage")
                || key.starts_with("b3") // Support B3 propagation as well
        })
        .collect();

    if env_vars.is_empty() {
        tracing::debug!("No trace context found in environment variables");
        return;
    }

    tracing::debug!(
        "Extracting trace context from {} environment variables: {:?}",
        env_vars.len(),
        env_vars.keys().collect::<Vec<_>>()
    );

    let extractor = EnvExtractor(&env_vars);
    let context =
        opentelemetry::global::get_text_map_propagator(|propagator| propagator.extract(&extractor));

    let span = tracing::Span::current();
    tracing::debug!("Setting parent context on span: {:?}", span.id());
    let _ = span.set_parent(context);
}

/// Extract trace context from a custom environment variable map and set it on the current span
///
/// This is useful when you want to extract trace context from a specific set of environment
/// variables rather than the current process environment.
///
/// # Arguments
///
/// * `env_vars` - A HashMap containing environment variables with potential trace context
///
/// # Examples
///
/// ```rust,ignore
/// use rialo_telemetry::extract_and_set_trace_context_from_env_map;
/// use HashMap;
///
/// let mut custom_env = HashMap::new();
/// custom_env.insert("traceparent".to_string(), "00-trace-id-span-id-01".to_string());
///
/// // Extract from custom environment map
/// extract_and_set_trace_context_from_env_map(&custom_env);
/// ```
pub fn extract_and_set_trace_context_from_env_map(env_vars: &HashMap<String, String>) {
    // Filter for trace context related variables
    let trace_vars: HashMap<String, String> = env_vars
        .iter()
        .filter(|(key, _)| {
            key.starts_with("traceparent")
                || key.starts_with("tracestate")
                || key.starts_with("baggage")
                || key.starts_with("b3")
        })
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();

    if trace_vars.is_empty() {
        tracing::debug!("No trace context found in provided environment variables");
        return;
    }

    tracing::debug!(
        "Extracting trace context from {} environment variables: {:?}",
        trace_vars.len(),
        trace_vars.keys().collect::<Vec<_>>()
    );

    let extractor = EnvExtractor(&trace_vars);
    let context =
        opentelemetry::global::get_text_map_propagator(|propagator| propagator.extract(&extractor));

    let span = tracing::Span::current();
    tracing::debug!("Setting parent context on span: {:?}", span.id());
    let _ = span.set_parent(context);
}

/// Inject trace context directly into a std::process::Command for convenient subprocess tracing
///
/// This is a convenience function that combines `inject_trace_env()` with applying the
/// environment variables to a Command. This provides an ergonomic way to pass trace context
/// to subprocess commands.
///
/// # Arguments
///
/// * `command` - The Command to inject trace context into
///
/// # Returns
///
/// The Command with trace context environment variables added
///
/// # Examples
///
/// ```rust,ignore
/// use rialo_telemetry::inject_trace_env_to_cmd;
/// use std::process::Command;
///
/// // In a traced context
/// let span = tracing::info_span!("subprocess_execution");
/// let _guard = span.enter();
///
/// // One-liner to create a command with trace context
/// let cmd = inject_trace_env_to_cmd(Command::new("worker"));
/// let output = cmd.arg("--task=process").output().expect("Failed to execute");
///
/// // Equivalent to:
/// // let trace_env = inject_trace_env();
/// // let mut cmd = Command::new("worker");
/// // for (key, value) in trace_env { cmd.env(key, value); }
/// ```
pub fn inject_trace_env_to_cmd(mut command: std::process::Command) -> std::process::Command {
    let trace_env = inject_trace_env();

    // Add all trace context environment variables to the command
    for (key, value) in trace_env {
        command.env(key, value);
    }

    command
}

// ==== Baggage Utilities ====
//
// Full baggage manipulation implementation using OpenTelemetry's Context API.
// All operations work by creating new contexts with modified baggage and setting
// them on the current span, which allows the baggage to be propagated through
// distributed tracing systems.

/// Get a baggage item from the current OpenTelemetry context.
///
/// # Arguments
///
/// * `key` - The baggage key to retrieve
///
/// # Returns
///
/// The baggage value if it exists, None otherwise
///
/// # Examples
///
/// ```rust,no_run
/// use rialo_telemetry::get_baggage;
/// use opentelemetry::{baggage::BaggageExt, Context};
///
/// // Set baggage using OpenTelemetry context
/// let baggage = opentelemetry::baggage::Baggage::builder()
///     .with_entry("user_id", "12345")
///     .build();
/// let cx = Context::current().with_baggage(baggage);
/// let _guard = cx.attach();
///
/// let user_id = get_baggage("user_id");
/// assert_eq!(user_id, Some("12345".to_string()));
/// # Ok::<(), anyhow::Error>(())
/// ```
pub fn get_baggage(key: &str) -> Option<String> {
    use opentelemetry::{baggage::BaggageExt, Context};

    Context::current().baggage().get(key).map(|v| v.to_string())
}

/// Get all baggage items from the current OpenTelemetry context.
///
/// # Returns
///
/// A HashMap containing all baggage key-value pairs
///
/// # Examples
///
/// ```rust,no_run
/// use rialo_telemetry::get_all_baggage;
/// use opentelemetry::{baggage::BaggageExt, Context};
///
/// // Set baggage using OpenTelemetry context
/// let baggage = opentelemetry::baggage::Baggage::builder()
///     .with_entry("user_id", "12345")
///     .with_entry("session_id", "session-abc")
///     .build();
/// let cx = Context::current().with_baggage(baggage);
/// let _guard = cx.attach();
///
/// let baggage = get_all_baggage();
/// assert_eq!(baggage.get("user_id"), Some(&"12345".to_string()));
/// # Ok::<(), anyhow::Error>(())
/// ```
pub fn get_all_baggage() -> HashMap<String, String> {
    use opentelemetry::{baggage::BaggageExt, Context};

    Context::current()
        .baggage()
        .into_iter()
        .map(|(k, (v, _metadata))| (k.to_string(), v.to_string()))
        .collect()
}

/// Parse protocol environment variable with default
fn parse_protocol_env(
    var_name: &str,
    default: crate::rialo_opentelemetry::Protocol,
) -> crate::rialo_opentelemetry::Protocol {
    std::env::var(var_name)
        .ok()
        .and_then(|v| match v.to_lowercase().as_str() {
            "grpc" => Some(crate::rialo_opentelemetry::Protocol::Grpc),
            "http/protobuf" => Some(crate::rialo_opentelemetry::Protocol::HttpProtobuf),
            "http/json" => Some(crate::rialo_opentelemetry::Protocol::HttpJson),
            _ => None,
        })
        .unwrap_or(default)
}

/// Parse sampling environment variable with default
fn parse_sampler_env(
    var_name: &str,
    arg_var_name: &str,
    default: crate::rialo_opentelemetry::Sampling,
) -> crate::rialo_opentelemetry::Sampling {
    std::env::var(var_name)
        .ok()
        .and_then(|v| match v.to_lowercase().as_str() {
            "always_on" | "parentbased_always_on" => {
                Some(crate::rialo_opentelemetry::Sampling::AlwaysOn)
            }
            "always_off" | "parentbased_always_off" => {
                Some(crate::rialo_opentelemetry::Sampling::AlwaysOff)
            }
            "traceidratio" | "parentbased_traceidratio" => {
                // Parse the ratio from OTEL_TRACES_SAMPLER_ARG
                let ratio = std::env::var(arg_var_name)
                    .ok()
                    .and_then(|arg| arg.parse::<f64>().ok())
                    .unwrap_or(0.05); // Default to 5% sampling if not specified or invalid

                // Clamp ratio to [0.0, 1.0] range
                let clamped_ratio = ratio.clamp(0.0, 1.0);
                Some(crate::rialo_opentelemetry::Sampling::Ratio(clamped_ratio))
            }
            _ => None,
        })
        .unwrap_or(default)
}

#[cfg(test)]
mod tests {
    use std::env;

    use serial_test::serial;

    use super::*;
    use crate::parse_env::parse_key_value_string;

    #[test]
    #[serial]
    fn test_otlp_config_default() {
        // Clear env vars to test defaults
        env::remove_var("OTEL_SERVICE_NAME");
        env::remove_var("OTEL_SERVICE_VERSION");
        env::remove_var("OTEL_EXPORTER_OTLP_ENDPOINT");
        env::remove_var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT");
        env::remove_var("OTEL_EXPORTER_OTLP_TRACES_HEADERS");
        env::remove_var("OTEL_EXPORTER_OTLP_HEADERS");
        env::remove_var("OTEL_TRACES_ENABLED");
        env::remove_var("OTEL_LOG_LEVEL");
        env::remove_var("OTEL_EXPORTER_OTLP_PROTOCOL");
        env::remove_var("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL");
        env::remove_var("OTEL_EXPORTER_OTLP_METRICS_PROTOCOL");
        env::remove_var("OTEL_RESOURCE_ATTRIBUTES");

        let config = OtlpConfig::default();
        assert_eq!(config.service_name, "rialo");
        assert_eq!(config.service_version, "unknown");
        assert_eq!(config.log_level, "info");
        assert!(config.traces_enabled);
        assert!(config.enable_console);
        // Should default to the same as general endpoint
        assert_eq!(
            config.effective_traces_endpoint(),
            Some(&"http://localhost:4318".to_string())
        );
        assert!(config.effective_traces_headers().is_empty());
        assert_eq!(config.exporter_protocol, Protocol::HttpProtobuf);
    }

    #[test]
    #[serial]
    fn test_otlp_config_from_env() {
        env::set_var("OTEL_SERVICE_NAME", "test-service");
        env::set_var("OTEL_SERVICE_VERSION", "1.2.3");
        env::set_var(
            "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT",
            "https://api.honeycomb.io/v1/traces",
        );
        env::set_var(
            "OTEL_EXPORTER_OTLP_TRACES_HEADERS",
            "x-honeycomb-team=your-api-key",
        );
        env::set_var("OTEL_TRACES_ENABLED", "true");
        env::set_var("OTEL_LOG_LEVEL", "debug");

        let config = OtlpConfig::default();
        assert_eq!(config.service_name, "test-service");
        assert_eq!(config.service_version, "1.2.3");
        assert_eq!(
            config.effective_traces_endpoint(),
            Some(&"https://api.honeycomb.io/v1/traces".to_string())
        );
        assert_eq!(
            config.traces_headers.get("x-honeycomb-team"),
            Some(&"your-api-key".to_string())
        );
        assert!(config.traces_enabled);
        assert_eq!(config.log_level, "debug");

        // Clean up
        env::remove_var("OTEL_SERVICE_NAME");
        env::remove_var("OTEL_SERVICE_VERSION");
        env::remove_var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT");
        env::remove_var("OTEL_EXPORTER_OTLP_TRACES_HEADERS");
        env::remove_var("OTEL_TRACES_ENABLED");
        env::remove_var("OTEL_LOG_LEVEL");
    }

    #[test]
    fn test_otlp_config_builder() {
        let config = OtlpConfig::new()
            .with_service_name("custom-service")
            .with_service_version("2.0.0")
            .with_traces_endpoint("https://api.honeycomb.io/v1/traces")
            .with_log_level("debug")
            .with_console_enabled(false)
            .with_traces_enabled(true);

        assert_eq!(config.service_name, "custom-service");
        assert_eq!(config.service_version, "2.0.0");
        assert_eq!(
            config.effective_traces_endpoint(),
            Some(&"https://api.honeycomb.io/v1/traces".to_string())
        );
        assert_eq!(config.log_level, "debug");
        assert!(!config.enable_console);
        assert!(config.traces_enabled);
    }

    #[test]
    fn test_parse_key_value_string() {
        let headers_str =
            "x-honeycomb-team=your-api-key,x-honeycomb-dataset=rialo,content-type=application/json";
        let headers = parse_key_value_string(headers_str);

        assert_eq!(headers.len(), 3);
        assert_eq!(
            headers.get("x-honeycomb-team"),
            Some(&"your-api-key".to_string())
        );
        assert_eq!(
            headers.get("x-honeycomb-dataset"),
            Some(&"rialo".to_string())
        );
        assert_eq!(
            headers.get("content-type"),
            Some(&"application/json".to_string())
        );
    }

    #[test]
    fn test_parse_key_value_string_semicolon() {
        let headers_str = "x-honeycomb-team=your-api-key;x-honeycomb-dataset=rialo";
        let headers = parse_key_value_string(headers_str);

        assert_eq!(headers.len(), 2);
        assert_eq!(
            headers.get("x-honeycomb-team"),
            Some(&"your-api-key".to_string())
        );
        assert_eq!(
            headers.get("x-honeycomb-dataset"),
            Some(&"rialo".to_string())
        );
    }

    #[test]
    fn test_parse_key_value_string_empty() {
        let headers = parse_key_value_string("");
        assert_eq!(headers.len(), 0);
    }

    #[test]
    #[serial]
    fn test_effective_endpoints() {
        // Clear any env vars that might interfere first
        env::remove_var("OTEL_EXPORTER_OTLP_ENDPOINT");
        env::remove_var("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT");

        let mut config = OtlpConfig::default();

        // Should have default endpoint same as general endpoint
        assert_eq!(
            config.effective_traces_endpoint(),
            Some(&"http://localhost:4318".to_string())
        );

        // Override general endpoint
        config.exporter_endpoint = Some("https://general.endpoint".to_string());
        assert_eq!(
            config.effective_traces_endpoint(),
            Some(&"https://general.endpoint".to_string())
        );

        // Traces-specific endpoint overrides general
        config.traces_exporter_endpoint = Some("https://traces.endpoint".to_string());
        assert_eq!(
            config.effective_traces_endpoint(),
            Some(&"https://traces.endpoint".to_string())
        );

        // Test with no general endpoint but with traces-specific
        config.exporter_endpoint = None;
        assert_eq!(
            config.effective_traces_endpoint(),
            Some(&"https://traces.endpoint".to_string())
        );

        // Test with neither set (should be None since we cleared exporter_endpoint)
        config.traces_exporter_endpoint = None;
        assert!(config.effective_traces_endpoint().is_none());
    }

    #[test]
    #[serial]
    fn test_effective_headers() {
        // Clear env vars to avoid interference
        env::remove_var("OTEL_EXPORTER_OTLP_HEADERS");
        env::remove_var("OTEL_EXPORTER_OTLP_TRACES_HEADERS");

        let mut config = OtlpConfig::default();

        // Set general headers
        config
            .headers
            .insert("general".to_string(), "value1".to_string());
        config
            .headers
            .insert("shared".to_string(), "general_value".to_string());

        // Set traces-specific headers
        config
            .traces_headers
            .insert("traces".to_string(), "value2".to_string());
        config
            .traces_headers
            .insert("shared".to_string(), "traces_value".to_string()); // Should override

        let effective = config.effective_traces_headers();
        assert_eq!(effective.len(), 3);
        assert_eq!(effective.get("general"), Some(&"value1".to_string()));
        assert_eq!(effective.get("traces"), Some(&"value2".to_string()));
        assert_eq!(effective.get("shared"), Some(&"traces_value".to_string())); // Traces-specific wins
    }

    #[tokio::test]
    #[serial]
    async fn test_inject_trace_headers() {
        use tracing::Instrument;

        // Initialize a basic tracer for testing
        let otlp_config = OtlpConfig::new()
            .with_service_name("test-service")
            .with_exporter_endpoint("".to_string()) // Empty to avoid network calls
            .with_traces_enabled(true);

        // This may fail if global subscriber is already set, which is fine in tests
        let _otel_result = init_otel(&otlp_config).await;

        // Create a test span with some context
        let span = tracing::info_span!("test_span", test_field = "test_value");

        let headers = async move {
            // Call inject_trace_headers inside the span
            inject_trace_headers()
        }
        .instrument(span)
        .await;

        // Verify we got some headers back (exact content depends on propagator setup)
        // The important thing is that the function works without panicking
        println!("Injected headers: {:?}", headers);

        // Verify headers structure - should be a valid HashMap
        // If there's an active trace, we should have traceparent header
        if !headers.is_empty() {
            // If headers are present, they should contain valid trace context keys
            let has_trace_headers = headers.contains_key("traceparent")
                || headers.contains_key("tracestate")
                || headers.keys().any(|k| k.starts_with("x-trace"));
            if !has_trace_headers {
                println!("Warning: Headers present but no recognized trace headers found");
            }
        }
    }

    #[test]
    fn test_inject_trace_headers_no_span() {
        // Test calling inject_trace_headers outside of any span
        // Should return empty headers or not panic
        let headers = inject_trace_headers();
        println!("Headers outside span: {:?}", headers);

        // Function should work even with no active span
        assert!(headers.is_empty() || !headers.is_empty());
    }

    #[tokio::test]
    #[serial]
    async fn test_inject_trace_env() {
        use tracing::Instrument;

        // Initialize a basic tracer for testing
        let otlp_config = OtlpConfig::new()
            .with_service_name("test-service")
            .with_exporter_endpoint("".to_string()) // Empty to avoid network calls
            .with_traces_enabled(true);

        // This may fail if global subscriber is already set, which is fine in tests
        let _otel_result = init_otel(&otlp_config).await;

        // Create a test span with some context
        let span = tracing::info_span!("test_env_span", test_field = "test_value");

        let env_vars = async move {
            // Call inject_trace_env inside the span
            inject_trace_env()
        }
        .instrument(span)
        .await;

        // Verify we got some environment variables back
        println!("Injected environment variables: {:?}", env_vars);

        // Verify environment variables structure
        // If there's an active trace, we should have trace context env vars
        if !env_vars.is_empty() {
            // Check for standard W3C trace context environment variables
            let has_trace_env =
                env_vars.contains_key("traceparent") || env_vars.contains_key("tracestate");
            if !has_trace_env {
                println!(
                    "Warning: Environment variables present but no recognized trace context found"
                );
            }

            // Validate traceparent format if present
            if let Some(traceparent) = env_vars.get("traceparent") {
                // Basic validation: should be in format "00-{trace-id}-{span-id}-{trace-flags}"
                let parts: Vec<&str> = traceparent.split('-').collect();
                assert_eq!(
                    parts.len(),
                    4,
                    "traceparent should have 4 parts separated by hyphens"
                );
                assert_eq!(parts[0], "00", "traceparent should start with version '00'");
            }
        }
    }

    #[test]
    fn test_inject_trace_env_no_span() {
        // Test calling inject_trace_env outside of any span
        let env_vars = inject_trace_env();
        println!("Environment variables outside span: {:?}", env_vars);

        // Function should work even with no active span
        assert!(env_vars.is_empty() || !env_vars.is_empty());
    }

    #[test]
    fn test_extract_and_set_trace_context_env_no_vars() {
        // Test extraction when no trace context environment variables are present
        // Should not panic and should handle gracefully
        extract_and_set_trace_context_env();

        // Function should complete without panicking when no trace context is found
    }

    #[test]
    fn test_extract_and_set_trace_context_from_env_map() {
        // Test with empty map
        let empty_env = HashMap::new();
        extract_and_set_trace_context_from_env_map(&empty_env);

        // Test with mock trace context data
        let mut mock_env = HashMap::new();
        mock_env.insert(
            "traceparent".to_string(),
            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string(),
        );
        mock_env.insert(
            "tracestate".to_string(),
            "rojo=00f067aa0ba902b7".to_string(),
        );
        mock_env.insert(
            "irrelevant_var".to_string(),
            "should_be_ignored".to_string(),
        );

        // Should not panic and should filter correctly
        extract_and_set_trace_context_from_env_map(&mock_env);
    }

    #[test]
    fn test_env_context_round_trip() {
        use tracing::Instrument;

        // This test simulates a parent process injecting context and a child extracting it
        let rt = tokio::runtime::Runtime::new().unwrap();

        rt.block_on(async {
            // Initialize telemetry
            let otlp_config = OtlpConfig::new()
                .with_service_name("test-parent")
                .with_exporter_endpoint("".to_string())
                .with_traces_enabled(true);

            let _otel_result = init_otel(&otlp_config).await;

            // Parent process: create a span and inject context
            let parent_span = tracing::info_span!("parent_process", process_id = "parent");

            let injected_env = async move { inject_trace_env() }
                .instrument(parent_span)
                .await;

            println!("Parent injected env vars: {:?}", injected_env);

            // Child process: extract context from the environment
            // (In real usage, this would be in a different process)
            let child_span = tracing::info_span!("child_process", process_id = "child");
            let _child_guard = child_span.enter();

            // Simulate child process extracting from parent's environment
            extract_and_set_trace_context_from_env_map(&injected_env);
        });
    }

    #[test]
    fn test_inject_trace_env_to_cmd() {
        use std::process::Command;

        // Test the ergonomic Command helper
        let cmd = inject_trace_env_to_cmd(Command::new("echo"));

        // We can't easily test the actual environment variables set on the Command,
        // but we can verify the function doesn't panic and returns a Command
        let program = cmd.get_program();
        assert_eq!(program, "echo");
    }

    #[tokio::test]
    #[serial]
    async fn test_inject_trace_env_to_cmd_with_span() {
        use std::process::Command;

        use tracing::Instrument;

        // Initialize basic tracer
        let otlp_config = OtlpConfig::new()
            .with_service_name("test-cmd")
            .with_exporter_endpoint("".to_string())
            .with_traces_enabled(true);

        let _otel_result = init_otel(&otlp_config).await;

        // Test within a span context
        let span = tracing::info_span!("cmd_test");
        let cmd = async move { inject_trace_env_to_cmd(Command::new("test-command")) }
            .instrument(span)
            .await;

        assert_eq!(cmd.get_program(), "test-command");
    }

    #[test]
    fn test_baggage_read_operations() {
        // Test the read-only baggage operations that work
        let value = get_baggage("non_existent_key");
        assert_eq!(value, None);

        let all_baggage = get_all_baggage();
        // Should return empty map when no baggage is set
        assert!(all_baggage.is_empty());
    }

    #[tokio::test]
    #[serial]
    async fn test_baggage_operations() {
        // Initialize OpenTelemetry for baggage to work properly
        let otlp_config = OtlpConfig::new()
            .with_service_name("test-baggage")
            .with_exporter_endpoint("".to_string())
            .with_traces_enabled(true);

        // This may fail if global subscriber is already set, which is fine in tests
        let otel_result = init_otel(&otlp_config).await;
        println!(
            "OpenTelemetry initialization result: {:?}",
            otel_result.is_ok()
        );
        if let Ok(result) = &otel_result {
            println!("Tracer available: {}", result.tracer.is_some());
        }

        use tracing::Instrument;

        // Create a span to provide context for baggage operations
        let test_span = tracing::info_span!("baggage_test");

        async {
            // Test basic baggage operations
            use opentelemetry::{
                baggage::{Baggage, BaggageExt, BaggageMetadata},
                Context, Key, Value,
            };

            // Create a context with baggage directly
            let mut baggage = Baggage::new();
            let _ = baggage.insert_with_metadata(
                Key::new("test_key".to_string()),
                Value::from("test_value".to_string()),
                BaggageMetadata::default(),
            );

            let context = Context::current().with_baggage(baggage);
            let _guard = context.attach();

            // Test that we can get baggage from the current context
            let current_context = Context::current();
            let baggage = current_context.baggage();
            println!("Direct context baggage: {:?}", baggage.get("test_key"));

            // Test our functions within the context scope
            {
                // Create a new context with our baggage
                let mut new_baggage = Baggage::new();
                let _ = new_baggage.insert_with_metadata(
                    Key::new("user_id".to_string()),
                    Value::from("12345".to_string()),
                    BaggageMetadata::default(),
                );
                let _ = new_baggage.insert_with_metadata(
                    Key::new("session_id".to_string()),
                    Value::from("session-abc".to_string()),
                    BaggageMetadata::default(),
                );

                let new_context = Context::current().with_baggage(new_baggage);
                let _inner_guard = new_context.attach();

                let user_id = get_baggage("user_id");
                let session_id = get_baggage("session_id");
                println!(
                    "After setting baggage in context, get_baggage('user_id') = {:?}",
                    user_id
                );
                println!(
                    "After setting baggage in context, get_baggage('session_id') = {:?}",
                    session_id
                );

                // Verify the values were set
                assert_eq!(user_id, Some("12345".to_string()));
                assert_eq!(session_id, Some("session-abc".to_string()));
                assert_eq!(get_baggage("non_existent"), None);

                // Test get_all_baggage
                let all_baggage = get_all_baggage();
                assert!(all_baggage.contains_key("user_id"));
                assert!(all_baggage.contains_key("session_id"));
                assert_eq!(all_baggage.get("user_id"), Some(&"12345".to_string()));
                assert_eq!(
                    all_baggage.get("session_id"),
                    Some(&"session-abc".to_string())
                );

                // Test removing baggage by manually creating a new context without the key
                {
                    let mut new_baggage = Baggage::new();
                    let _ = new_baggage.insert_with_metadata(
                        Key::new("user_id".to_string()),
                        Value::from("12345".to_string()),
                        BaggageMetadata::default(),
                    );
                    // Note: session_id is not added, so it's effectively removed

                    let new_context = Context::current().with_baggage(new_baggage);
                    let _inner_guard = new_context.attach();

                    assert_eq!(get_baggage("session_id"), None);
                    assert_eq!(get_baggage("user_id"), Some("12345".to_string()));
                    // Should still exist
                }

                // Test clearing baggage by creating an empty context
                {
                    let empty_baggage = Baggage::new();
                    let new_context = Context::current().with_baggage(empty_baggage);
                    let _inner_guard = new_context.attach();

                    assert_eq!(get_baggage("user_id"), None);
                    assert!(get_all_baggage().is_empty());
                }
            }
        }
        .instrument(test_span)
        .await;
    }

    /// Regression test for tracing-opentelemetry issue #183: panic in
    /// `on_follows_from` when a followed span is being closed on another thread.
    ///
    /// In tracing-opentelemetry <=0.31.0, the `on_follows_from` handler called
    /// `.expect()` on OTel extension data that could already be removed by a
    /// concurrent `on_close` on another thread.  This caused the panic:
    ///   "Missing otel data span extensions"
    ///
    /// The fix (0.32.x) replaces `.expect()` with a `let-else` that gracefully
    /// ignores missing extension data.
    ///
    /// This test deterministically reproduces the race by using a `WaitOnCloseLayer`
    /// that sleeps during `on_close`, widening the window for `follows_from` to
    /// observe the partially-closed state. Adapted from the upstream fix's test
    /// (tokio-rs/tracing-opentelemetry PR #244).
    #[test]
    fn test_follows_from_during_close_no_panic() {
        use opentelemetry::trace::TracerProvider as _;
        use tracing::Subscriber;
        use tracing_subscriber::prelude::*;

        /// Layer that artificially delays `on_close` processing, creating a
        /// window where the span's OTel extension data has been removed but the
        /// span still exists in the Registry.
        #[derive(Clone)]
        struct WaitOnCloseLayer;
        impl<S: Subscriber> tracing_subscriber::Layer<S> for WaitOnCloseLayer {
            fn on_close(&self, _id: tracing::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
                std::thread::sleep(std::time::Duration::from_millis(20));
            }
        }

        // In-memory exporter (no network required)
        #[derive(Clone, Default, Debug)]
        struct NoopExporter;
        impl opentelemetry_sdk::trace::SpanExporter for NoopExporter {
            async fn export(
                &self,
                _batch: Vec<opentelemetry_sdk::trace::SpanData>,
            ) -> opentelemetry_sdk::error::OTelSdkResult {
                Ok(())
            }
        }

        let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
            .with_simple_exporter(NoopExporter)
            .build();
        let tracer = provider.tracer("regression-test");

        // Stack: Registry → OpenTelemetryLayer → WaitOnCloseLayer
        // The OTel layer removes extension data in on_close, then
        // WaitOnCloseLayer sleeps 20ms keeping the span alive in the Registry.
        let subscriber = std::sync::Arc::new(
            tracing_subscriber::registry()
                .with(tracing_opentelemetry::layer().with_tracer(tracer))
                .with(WaitOnCloseLayer),
        );

        tracing::subscriber::with_default(subscriber, || {
            let cause_span = tracing::debug_span!("cause");
            let cause_id = cause_span.id().unwrap();

            std::thread::scope(|scope| {
                // Thread 1: drops the cause span, triggering on_close which
                // removes OTel extension data then sleeps 20ms.
                scope.spawn(|| {
                    drop(cause_span);
                });

                // Wait 5ms so we're inside the WaitOnCloseLayer sleep window:
                // OTel data is gone but the span still exists in the Registry.
                std::thread::sleep(std::time::Duration::from_millis(5));

                // Thread 0: calls follows_from on the partially-closed span.
                // In <=0.31.0 this panics with "Missing otel data span extensions".
                // In >=0.32.0 this gracefully handles the missing data.
                let effect_span = tracing::debug_span!("effect");
                effect_span.follows_from(cause_id);
            });
        });

        let _ = provider.shutdown();
    }
}