//! OpenTelemetry distributed tracing with OTLP export.
//!
//! This module provides distributed tracing capabilities using OpenTelemetry
//! with OTLP (OpenTelemetry Protocol) export. It enables end-to-end request
//! tracking across service boundaries with parent-child span relationships
//! and context propagation.
//!
//! # Features
//!
//! - Automatic span creation and linking across service boundaries
//! - Parent-child span relationships via gRPC metadata
//! - Trace context propagation (trace_id, span_id, parent_span_id)
//! - OTLP export (compatible with Jaeger, Grafana Tempo, Datadog, etc.)
//! - Integration with correlation IDs
//!
//! # Architecture
//!
//! The tracing system uses OpenTelemetry's trace API to create spans that
//! represent units of work. Spans are linked into traces that follow requests
//! across multiple services. Trace context is propagated via gRPC metadata
//! using the W3C Trace Context specification.

use crate::error::{CommonError, CommonResult};
use opentelemetry::{global, trace::TracerProvider as _, KeyValue};
use opentelemetry_otlp::{SpanExporter, WithTonicConfig};
use opentelemetry_sdk::{
    trace::{Sampler, SdkTracerProvider},
    Resource,
};
use std::sync::OnceLock;

/// Dedicated tokio runtime for the OTLP batch exporter in sync binaries.
///
/// The batch span processor uses `tokio::spawn` internally for periodic flushing
/// and the tonic gRPC transport needs a reactor. Async services (with
/// `#[tokio::main]`) already provide this; sync training binaries need a
/// dedicated runtime that lives for the process lifetime.
///
/// Uses a multi-thread runtime with 1 worker thread so the export loop is
/// actively polled without the caller needing to drive it.
static OTEL_RUNTIME: OnceLock<Option<tokio::runtime::Runtime>> = OnceLock::new();

/// Configuration for OpenTelemetry tracing.
///
/// Controls all aspects of distributed tracing including service identification,
/// export endpoints, and sampling strategy. Uses OTLP (OpenTelemetry Protocol)
/// for trace export, which is the vendor-neutral standard supported by Jaeger,
/// Tempo, Datadog, and other backends.
#[derive(Debug, Clone)]
pub struct TracingConfig {
    /// Service name to identify in traces
    pub service_name: String,

    /// OTLP collector endpoint (e.g., "http://localhost:4317" for gRPC)
    pub otlp_endpoint: String,

    /// Enable OTLP export (disable for testing)
    pub enable_export: bool,
}

impl Default for TracingConfig {
    fn default() -> Self {
        Self {
            service_name: "foxhunt".to_string(),
            otlp_endpoint: "http://localhost:4317".to_string(),
            enable_export: true,
        }
    }
}

/// Build an OpenTelemetry tracer backed by an OTLP exporter.
///
/// Returns `Some(tracer)` when OTLP export is enabled and the exporter can be
/// constructed, or `None` on failure (graceful degradation). The caller wraps
/// the tracer in a `tracing_opentelemetry::layer().with_tracer(tracer)` layer
/// inside their subscriber stack so the `S` type parameter is inferred
/// correctly.
///
/// This function also registers the tracer provider as the global provider
/// (required for propagation helpers like `inject_trace_context`).
///
/// # Arguments
///
/// * `config` - Tracing configuration
///
/// # Returns
///
/// Returns `Some(Tracer)` on success, or `None` if export is disabled or the
/// OTLP exporter fails to build.
pub fn build_otel_tracer(
    config: &TracingConfig,
) -> Option<opentelemetry_sdk::trace::Tracer> {
    if !config.enable_export {
        return None;
    }

    // The batch exporter needs a tokio runtime for gRPC transport and periodic
    // flush. Async services already have one; sync training binaries get a
    // dedicated single-worker runtime that lives for the process lifetime.
    let _runtime_guard = if tokio::runtime::Handle::try_current().is_ok() {
        None
    } else {
        let rt_opt = OTEL_RUNTIME.get_or_init(|| {
            tokio::runtime::Builder::new_multi_thread()
                .worker_threads(1)
                .thread_name("otel-export")
                .enable_all()
                .build()
                .ok()
        });
        match rt_opt.as_ref() {
            Some(rt) => Some(rt.enter()),
            None => {
                eprintln!("OTLP runtime creation failed (non-fatal)");
                return None;
            }
        }
    };

    // Build an explicit tonic Channel that respects the URI scheme (http vs https).
    // This prevents TLS negotiation on plaintext http:// endpoints like Tempo.
    let channel = match tonic::transport::Channel::from_shared(config.otlp_endpoint.clone()) {
        Ok(endpoint) => endpoint.connect_lazy(),
        Err(e) => {
            eprintln!("Invalid OTLP endpoint URI (non-fatal): {e}");
            return None;
        }
    };

    let exporter = match SpanExporter::builder()
        .with_tonic()
        .with_channel(channel)
        .build()
    {
        Ok(e) => e,
        Err(e) => {
            eprintln!("OTLP exporter build failed (non-fatal): {e}");
            return None;
        }
    };

    let provider = SdkTracerProvider::builder()
        .with_sampler(Sampler::TraceIdRatioBased(0.1))
        .with_batch_exporter(exporter)
        .with_resource(
            Resource::builder()
                .with_service_name(config.service_name.clone())
                .with_attribute(KeyValue::new("service.version", env!("CARGO_PKG_VERSION")))
                .build(),
        )
        .build();

    let tracer = provider.tracer(config.service_name.clone());
    global::set_tracer_provider(provider);

    Some(tracer)
}

/// Initialize OpenTelemetry tracing with OTLP export.
///
/// This function sets up the OpenTelemetry tracing pipeline with OTLP
/// (OpenTelemetry Protocol) as the trace exporter. OTLP is the vendor-neutral
/// standard supported by Jaeger (v1.35+), Grafana Tempo, Datadog, and others.
///
/// **Note:** Prefer [`build_otel_tracer`] for new code. This function is kept
/// for backward compatibility but only sets the global tracer provider without
/// adding a layer to the subscriber.
///
/// # Arguments
///
/// * `config` - Tracing configuration
///
/// # Returns
///
/// Returns `Ok(())` on success, or an error if initialization fails.
///
/// # Example
///
/// ```rust,no_run
/// use common::observability::tracing_config::{init_tracing, TracingConfig};
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     let config = TracingConfig {
///         service_name: "trading_service".to_string(),
///         otlp_endpoint: "http://localhost:4317".to_string(),
///         enable_export: true,
///     };
///
///     init_tracing(config).await?;
///
///     // Service logic with automatic tracing
///     Ok(())
/// }
/// ```
pub async fn init_tracing(config: TracingConfig) -> CommonResult<()> {
    if !config.enable_export {
        // Tracing disabled for testing - use noop tracer
        return Ok(());
    }

    // Build an explicit tonic Channel — see build_otel_tracer() comment for rationale
    let channel = tonic::transport::Channel::from_shared(config.otlp_endpoint.clone())
        .map_err(|e| CommonError::config(format!("Invalid OTLP endpoint URI: {}", e)))?
        .connect_lazy();

    let exporter = SpanExporter::builder()
        .with_tonic()
        .with_channel(channel)
        .build()
        .map_err(|e| CommonError::config(format!("Failed to build OTLP exporter: {}", e)))?;

    // Create a tracer provider with the OTLP exporter using batched export
    let provider = SdkTracerProvider::builder()
        .with_sampler(Sampler::TraceIdRatioBased(0.1))
        .with_batch_exporter(exporter)
        .with_resource(
            Resource::builder()
                .with_service_name(config.service_name.clone())
                .with_attribute(KeyValue::new("service.version", env!("CARGO_PKG_VERSION")))
                .build(),
        )
        .build();

    // Set as global tracer provider
    global::set_tracer_provider(provider);

    Ok(())
}

/// Shutdown the OpenTelemetry tracing pipeline gracefully.
///
/// This function should be called during service shutdown to ensure all
/// pending spans are flushed to Jaeger before the process exits.
///
/// # Example
///
/// ```rust,no_run
/// use common::observability::tracing_config::shutdown_tracing;
///
/// async fn shutdown() {
///     shutdown_tracing().await;
/// }
/// ```
pub async fn shutdown_tracing() {
    // In opentelemetry 0.31, shutdown is on the provider itself via Drop.
    // The global provider is shut down when it goes out of scope.
    // Explicit flush is no longer needed via a global function.
}

/// Extract trace context from gRPC metadata.
///
/// This function extracts the W3C Trace Context headers from gRPC metadata
/// and returns them as a tuple of (trace_id, span_id, parent_span_id).
///
/// # Arguments
///
/// * `metadata` - gRPC metadata map
///
/// # Returns
///
/// Returns `Some((trace_id, span_id, parent_span_id))` if trace context is
/// present, or `None` if not found.
pub fn extract_trace_context(
    metadata: &tonic::metadata::MetadataMap,
) -> Option<(String, String, Option<String>)> {
    let traceparent = metadata
        .get("traceparent")
        .and_then(|v| v.to_str().ok())?;

    // Parse W3C Trace Context format: "00-{trace_id}-{span_id}-{flags}"
    let parts: Vec<&str> = traceparent.split('-').collect();
    if parts.len() != 4 {
        return None;
    }

    let trace_id = parts[1].to_string();
    let parent_span_id = parts[2].to_string();

    // Generate new span ID for this service
    let span_id = format!("{:016x}", rand::random::<u64>());

    Some((trace_id, span_id, Some(parent_span_id)))
}

/// Inject trace context into gRPC metadata.
///
/// This function creates W3C Trace Context headers and injects them into
/// gRPC metadata for propagation to downstream services.
///
/// # Arguments
///
/// * `metadata` - Mutable gRPC metadata map
/// * `trace_id` - Trace ID
/// * `span_id` - Current span ID
///
/// # Example
///
/// ```rust,no_run
/// use common::observability::tracing_config::inject_trace_context;
/// use tonic::metadata::MetadataMap;
///
/// let mut metadata = MetadataMap::new();
/// inject_trace_context(&mut metadata, "trace123", "span456");
/// ```
pub fn inject_trace_context(
    metadata: &mut tonic::metadata::MetadataMap,
    trace_id: &str,
    span_id: &str,
) {
    // Create W3C Trace Context header: "00-{trace_id}-{span_id}-01"
    let traceparent = format!("00-{}-{}-01", trace_id, span_id);

    if let Ok(value) = tonic::metadata::MetadataValue::try_from(&traceparent) {
        metadata.insert("traceparent", value);
    }
}

/// Create a new trace context for a root span.
///
/// This function generates a new trace ID and span ID for requests that
/// originate at the API Gateway (i.e., no parent trace context exists).
///
/// # Returns
///
/// Returns a tuple of (trace_id, span_id).
///
/// # Example
///
/// ```rust
/// use common::observability::tracing_config::create_root_trace_context;
///
/// let (trace_id, span_id) = create_root_trace_context();
/// println!("Trace: {}, Span: {}", trace_id, span_id);
/// ```
pub fn create_root_trace_context() -> (String, String) {
    let trace_id = format!("{:032x}", rand::random::<u128>());
    let span_id = format!("{:016x}", rand::random::<u64>());
    (trace_id, span_id)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tonic::metadata::MetadataMap;

    #[test]
    fn test_create_root_trace_context() {
        let (trace_id1, span_id1) = create_root_trace_context();
        let (trace_id2, span_id2) = create_root_trace_context();

        // IDs should be unique
        assert_ne!(trace_id1, trace_id2);
        assert_ne!(span_id1, span_id2);

        // Trace ID should be 32 hex characters
        assert_eq!(trace_id1.len(), 32);
        assert!(trace_id1.chars().all(|c| c.is_ascii_hexdigit()));

        // Span ID should be 16 hex characters
        assert_eq!(span_id1.len(), 16);
        assert!(span_id1.chars().all(|c| c.is_ascii_hexdigit()));
    }

    #[test]
    fn test_inject_and_extract_trace_context() {
        let mut metadata = MetadataMap::new();
        let trace_id = "0123456789abcdef0123456789abcdef";
        let span_id = "0123456789abcdef";

        // Inject trace context
        inject_trace_context(&mut metadata, trace_id, span_id);

        // Verify traceparent header is set
        let traceparent = metadata.get("traceparent");
        assert!(traceparent.is_some());

        let traceparent_str = traceparent.unwrap().to_str().unwrap();
        assert_eq!(
            traceparent_str,
            format!("00-{}-{}-01", trace_id, span_id)
        );

        // Extract trace context
        let extracted = extract_trace_context(&metadata);
        assert!(extracted.is_some());

        let (extracted_trace_id, extracted_span_id, parent_span_id) = extracted.unwrap();
        assert_eq!(extracted_trace_id, trace_id);
        assert_ne!(extracted_span_id, span_id); // New span ID generated
        assert_eq!(parent_span_id, Some(span_id.to_string()));
    }

    #[test]
    fn test_extract_trace_context_missing() {
        let metadata = MetadataMap::new();

        // Should return None when traceparent is missing
        let extracted = extract_trace_context(&metadata);
        assert!(extracted.is_none());
    }

    #[test]
    fn test_extract_trace_context_invalid_format() {
        let mut metadata = MetadataMap::new();

        // Invalid traceparent format
        let value = tonic::metadata::MetadataValue::try_from("invalid-format").unwrap();
        metadata.insert("traceparent", value);

        let extracted = extract_trace_context(&metadata);
        assert!(extracted.is_none());
    }

    #[tokio::test]
    async fn test_init_tracing_disabled() {
        let config = TracingConfig {
            service_name: "test_service".to_string(),
            otlp_endpoint: "http://localhost:4317".to_string(),
            enable_export: false,
        };

        // Should succeed when disabled
        let result = init_tracing(config).await;
        assert!(result.is_ok());
    }

    #[test]
    fn test_default_config() {
        let config = TracingConfig::default();

        assert_eq!(config.service_name, "foxhunt");
        assert_eq!(config.otlp_endpoint, "http://localhost:4317");
        assert!(config.enable_export);
    }
}
