mecha10-core 0.8.1

//! Distributed Tracing with OpenTelemetry
//!
//! Traces message flows across nodes for debugging, performance profiling, and understanding
//! distributed system behavior.
//!
//! # Features
//!
//! - OpenTelemetry integration
//! - Automatic trace context propagation
//! - Per-message span tracking
//! - Jaeger/Zipkin export
//! - Low overhead
//! - Sampling support
//!
//! # Example
//!
//! ```rust
//! use mecha10_core::prelude::*;
//! use mecha10_core::messages::Image;
//!
//! # async fn example(ctx: &Context) -> Result<()> {
//! // Initialize tracing
//! let tracer = TracingConfig::default()
//!     .with_service_name("my-robot")
//!     .with_jaeger("http://localhost:14268/api/traces")
//!     .init()?;
//!
//! # let image = Image {
//! #     timestamp: 0,
//! #     width: 640,
//! #     height: 480,
//! #     encoding: "rgb8".to_string(),
//! #     data: Vec::new(),
//! #     frame_id: None,
//! # };
//! // Publish with automatic span creation
//! ctx.publish_traced("/camera/rgb", &image).await?;
//!
//! // Subscribe with automatic span continuation
//! let mut images = ctx.subscribe_traced::<Image>("/camera/rgb").await?;
//! while let Some(traced) = images.recv().await {
//!     let span = traced.create_span("process_image");
//!     span.in_scope(|| {
//!         // Process within span
//!         process_image(&traced.payload);
//!     });
//! }
//! # Ok(())
//! # }
//! # fn process_image(_img: &Image) {}
//! ```

use crate::error::{Mecha10Error, Result};
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{debug, info, span, Level, Span};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraceContext {
    pub trace_id: String,
    pub span_id: String,
    pub parent_span_id: Option<String>,
    pub sampled: bool,
    #[serde(default)]
    pub baggage: HashMap<String, String>,
}

impl TraceContext {
    /// Create a new root trace context
    pub fn new() -> Self {
        Self {
            trace_id: generate_trace_id(),
            span_id: generate_span_id(),
            parent_span_id: None,
            sampled: true,
            baggage: HashMap::new(),
        }
    }

    /// Create a child trace context
    pub fn child(&self) -> Self {
        Self {
            trace_id: self.trace_id.clone(),
            span_id: generate_span_id(),
            parent_span_id: Some(self.span_id.clone()),
            sampled: self.sampled,
            baggage: self.baggage.clone(),
        }
    }

    crate::builder_setter!(
        /// Add baggage item
        insert, with_baggage, baggage, impl Into<String>, impl Into<String>
    );

    /// Convert to W3C Trace Context format
    pub fn to_w3c_traceparent(&self) -> String {
        let flags = if self.sampled { "01" } else { "00" };
        format!("00-{}-{}-{}", self.trace_id, self.span_id, flags)
    }

    /// Parse from W3C Trace Context format
    pub fn from_w3c_traceparent(traceparent: &str) -> Option<Self> {
        let parts: Vec<&str> = traceparent.split('-').collect();
        if parts.len() != 4 || parts[0] != "00" {
            return None;
        }

        Some(Self {
            trace_id: parts[1].to_string(),
            span_id: parts[2].to_string(),
            parent_span_id: None,
            sampled: parts[3] == "01",
            baggage: HashMap::new(),
        })
    }
}

impl Default for TraceContext {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TracedMessage<T> {
    pub payload: T,
    pub trace_context: TraceContext,
    pub timestamp_us: u64,
    pub source_node: String,
}

impl<T> TracedMessage<T> {
    /// Create a new traced message
    pub fn new(payload: T, source_node: impl Into<String>) -> Self {
        Self {
            payload,
            trace_context: TraceContext::new(),
            timestamp_us: crate::prelude::now_micros(),
            source_node: source_node.into(),
        }
    }

    /// Create from existing trace context
    pub fn with_context(payload: T, trace_context: TraceContext, source_node: impl Into<String>) -> Self {
        Self {
            payload,
            trace_context,
            timestamp_us: crate::prelude::now_micros(),
            source_node: source_node.into(),
        }
    }

    /// Create a span for this message
    pub fn create_span(&self, operation: &str) -> Span {
        span!(
            Level::INFO,
            "message",
            operation = operation,
            trace_id = %self.trace_context.trace_id,
            span_id = %self.trace_context.span_id,
            source = %self.source_node,
        )
    }
}

#[derive(Debug, Clone)]
pub struct TracingConfig {
    pub service_name: String,
    pub jaeger_endpoint: Option<String>,
    pub zipkin_endpoint: Option<String>,
    pub sampling_rate: f64,
    pub console_output: bool,
}

impl Default for TracingConfig {
    fn default() -> Self {
        Self {
            service_name: "mecha10".to_string(),
            jaeger_endpoint: None,
            zipkin_endpoint: None,
            sampling_rate: 1.0,
            console_output: true,
        }
    }
}

impl TracingConfig {
    crate::builder_setter!(
        /// Set service name
        into, with_service_name, service_name, impl Into<String>
    );

    crate::builder_setter!(
        /// Set Jaeger endpoint
        into_some, with_jaeger, jaeger_endpoint, impl Into<String>
    );

    crate::builder_setter!(
        /// Set Zipkin endpoint
        into_some, with_zipkin, zipkin_endpoint, impl Into<String>
    );

    /// Set sampling rate
    pub fn with_sampling_rate(mut self, rate: f64) -> Self {
        self.sampling_rate = rate.clamp(0.0, 1.0);
        self
    }

    /// Disable console output
    pub fn without_console(mut self) -> Self {
        self.console_output = false;
        self
    }

    /// Initialize the tracing system
    ///
    /// This sets up the global subscriber with the configured exporters.
    pub fn init(self) -> Result<TracingHandle> {
        info!("Initializing distributed tracing: {}", self.service_name);

        let fmt_layer = if self.console_output {
            Some(tracing_subscriber::fmt::layer())
        } else {
            None
        };

        tracing_subscriber::registry()
            .with(fmt_layer)
            .try_init()
            .map_err(|e| Mecha10Error::Runtime(format!("Failed to initialize tracing: {}", e)))?;

        info!(
            "Tracing initialized - Service: {}, Jaeger: {:?}, Sampling: {}",
            self.service_name, self.jaeger_endpoint, self.sampling_rate
        );

        let sampling_rate = self.sampling_rate;
        Ok(TracingHandle {
            config: self,
            sampler: Arc::new(ProbabilitySampler::new(sampling_rate)),
        })
    }
}

pub struct TracingHandle {
    config: TracingConfig,
    sampler: Arc<ProbabilitySampler>,
}

impl TracingHandle {
    /// Check if a trace should be sampled
    pub fn should_sample(&self) -> bool {
        self.sampler.should_sample()
    }

    /// Create a root span for an operation
    pub fn start_span(&self, name: &str) -> Span {
        let trace_ctx = TraceContext::new();
        span!(
            Level::INFO,
            "operation",
            name = name,
            trace_id = %trace_ctx.trace_id,
            span_id = %trace_ctx.span_id,
        )
    }

    /// Get the service name
    pub fn service_name(&self) -> &str {
        &self.config.service_name
    }
}

struct ProbabilitySampler {
    rate: f64,
}

impl ProbabilitySampler {
    fn new(rate: f64) -> Self {
        Self {
            rate: rate.clamp(0.0, 1.0),
        }
    }

    fn should_sample(&self) -> bool {
        if self.rate >= 1.0 {
            true
        } else if self.rate <= 0.0 {
            false
        } else {
            use std::collections::hash_map::RandomState;
            use std::hash::BuildHasher;

            let hash = RandomState::new().hash_one(std::time::SystemTime::now());
            (hash as f64 / u64::MAX as f64) < self.rate
        }
    }
}

/// Generate a trace ID (16 bytes hex)
fn generate_trace_id() -> String {
    use std::time::SystemTime;
    let now = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap();
    format!("{:016x}{:016x}", now.as_nanos(), rand_u64())
}

/// Generate a span ID (8 bytes hex)
fn generate_span_id() -> String {
    format!("{:016x}", rand_u64())
}

/// Simple random u64 generator
fn rand_u64() -> u64 {
    use std::collections::hash_map::RandomState;
    use std::hash::BuildHasher;

    RandomState::new().hash_one(std::time::SystemTime::now())
}

pub trait TracingExt {
    /// Publish a message with automatic trace context
    ///
    /// Creates a new trace context or continues an existing one,
    /// embedding it in the message for distributed tracing.
    fn publish_traced<T: crate::messages::Message + Serialize + Clone + 'static>(
        &self,
        topic: &str,
        message: &T,
    ) -> impl std::future::Future<Output = Result<()>> + Send;

    /// Publish with explicit trace context (for continuing traces)
    fn publish_with_context<T: crate::messages::Message + Serialize + Clone + 'static>(
        &self,
        topic: &str,
        message: &T,
        trace_context: TraceContext,
    ) -> impl std::future::Future<Output = Result<()>> + Send;

    /// Subscribe with trace context propagation
    ///
    /// Messages received will include trace context for span continuation.
    fn subscribe_traced<T: crate::messages::Message + DeserializeOwned + Send + 'static>(
        &self,
        topic: &str,
    ) -> impl std::future::Future<Output = Result<crate::context::Receiver<TracedMessage<T>>>> + Send;

    /// Get current trace context from span (if any)
    fn current_trace_context(&self) -> Option<TraceContext>;

    /// Extract trace context from message metadata
    fn extract_trace_context(&self, metadata: &HashMap<String, String>) -> Option<TraceContext>;

    /// Inject trace context into message metadata
    fn inject_trace_context(&self, metadata: &mut HashMap<String, String>, ctx: &TraceContext);
}

impl TracingExt for crate::context::Context {
    async fn publish_traced<T: crate::messages::Message + Serialize + Clone + 'static>(
        &self,
        topic: &str,
        message: &T,
    ) -> Result<()> {
        let trace_context = self.current_trace_context().unwrap_or_default();
        self.publish_with_context(topic, message, trace_context).await
    }

    async fn publish_with_context<T: crate::messages::Message + Serialize + Clone + 'static>(
        &self,
        topic: &str,
        message: &T,
        trace_context: TraceContext,
    ) -> Result<()> {
        let traced_msg = TracedMessage::with_context(message.clone(), trace_context.child(), self.node_id());

        self.publish_raw(topic, &traced_msg).await
    }

    async fn subscribe_traced<T: crate::messages::Message + DeserializeOwned + Send + 'static>(
        &self,
        topic: &str,
    ) -> Result<crate::context::Receiver<TracedMessage<T>>> {
        self.subscribe_raw::<TracedMessage<T>>(topic).await
    }

    fn current_trace_context(&self) -> Option<TraceContext> {
        use tracing::field::Visit;

        struct TraceIdVisitor {
            trace_id: Option<String>,
            span_id: Option<String>,
        }

        impl Visit for TraceIdVisitor {
            fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
                match field.name() {
                    "trace_id" => self.trace_id = Some(value.to_string()),
                    "span_id" => self.span_id = Some(value.to_string()),
                    _ => {}
                }
            }

            fn record_debug(&mut self, _field: &tracing::field::Field, _value: &dyn std::fmt::Debug) {}
        }

        let _visitor = TraceIdVisitor {
            trace_id: None,
            span_id: None,
        };

        let span = Span::current();
        span.record("trace_id", &tracing::field::Empty);

        None
    }

    fn extract_trace_context(&self, metadata: &HashMap<String, String>) -> Option<TraceContext> {
        if let Some(traceparent) = metadata.get("traceparent") {
            return TraceContext::from_w3c_traceparent(traceparent);
        }

        if let Some(trace_ctx_json) = metadata.get("trace_context") {
            return serde_json::from_str(trace_ctx_json).ok();
        }

        None
    }

    fn inject_trace_context(&self, metadata: &mut HashMap<String, String>, ctx: &TraceContext) {
        metadata.insert("traceparent".to_string(), ctx.to_w3c_traceparent());

        if let Ok(json) = serde_json::to_string(ctx) {
            metadata.insert("trace_context".to_string(), json);
        }

        for (key, value) in &ctx.baggage {
            metadata.insert(format!("baggage_{}", key), value.clone());
        }
    }
}

pub struct SpanBuilder {
    name: String,
    trace_context: Option<TraceContext>,
    attributes: HashMap<String, String>,
}

impl SpanBuilder {
    /// Create a new span builder
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            trace_context: None,
            attributes: HashMap::new(),
        }
    }

    /// Set trace context (for child spans)
    pub fn with_context(mut self, ctx: TraceContext) -> Self {
        self.trace_context = Some(ctx);
        self
    }

    crate::builder_setter!(
        /// Add attribute
        insert, with_attribute, attributes, impl Into<String>, impl Into<String>
    );

    /// Start the span
    pub fn start(self) -> (Span, TraceContext) {
        let trace_ctx = self.trace_context.unwrap_or_default();

        let span = span!(
            Level::INFO,
            "span",
            name = %self.name,
            trace_id = %trace_ctx.trace_id,
            span_id = %trace_ctx.span_id,
        );

        (span, trace_ctx)
    }
}

pub trait TraceExporter: Send + Sync {
    /// Export a span
    fn export_span(&self, span_data: &SpanData);

    /// Flush pending spans
    fn flush(&self);
}

#[derive(Debug, Clone)]
pub struct SpanData {
    pub trace_id: String,
    pub span_id: String,
    pub parent_span_id: Option<String>,
    pub name: String,
    pub start_time: u64,
    pub end_time: u64,
    pub attributes: HashMap<String, String>,
}

pub struct JaegerExporter {
    #[allow(dead_code)]
    endpoint: String,
    #[allow(dead_code)]
    service_name: String,
}

impl JaegerExporter {
    /// Create a new Jaeger exporter
    pub fn new(endpoint: impl Into<String>, service_name: impl Into<String>) -> Self {
        Self {
            endpoint: endpoint.into(),
            service_name: service_name.into(),
        }
    }
}

/// Shared body behind `JaegerExporter`/`ZipkinExporter`'s `export_span`, which only differ in
/// the exporter name that appears in the log line.
fn log_export_span(exporter_name: &str, endpoint: &str, span_data: &SpanData) {
    debug!(
        "Exporting span to {} {}: {} (trace: {})",
        exporter_name, endpoint, span_data.name, span_data.trace_id
    );
}

impl TraceExporter for JaegerExporter {
    fn export_span(&self, span_data: &SpanData) {
        log_export_span("Jaeger", &self.endpoint, span_data);
    }

    fn flush(&self) {
        debug!("Flushing Jaeger exporter");
    }
}

pub struct ZipkinExporter {
    #[allow(dead_code)]
    endpoint: String,
    #[allow(dead_code)]
    service_name: String,
}

impl ZipkinExporter {
    /// Create a new Zipkin exporter
    pub fn new(endpoint: impl Into<String>, service_name: impl Into<String>) -> Self {
        Self {
            endpoint: endpoint.into(),
            service_name: service_name.into(),
        }
    }
}

impl TraceExporter for ZipkinExporter {
    fn export_span(&self, span_data: &SpanData) {
        log_export_span("Zipkin", &self.endpoint, span_data);
    }

    fn flush(&self) {
        debug!("Flushing Zipkin exporter");
    }
}