halldyll-core 0.1.0

//! Traces - Distributed tracing

use std::collections::HashMap;
use uuid::Uuid;

/// Trace context
#[derive(Debug, Clone)]
pub struct TraceContext {
    /// Trace ID
    pub trace_id: String,
    /// Span ID
    pub span_id: String,
    /// Parent Span ID
    pub parent_span_id: Option<String>,
    /// Attributes
    pub attributes: HashMap<String, String>,
}

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

impl TraceContext {
    /// New trace context
    pub fn new() -> Self {
        Self {
            trace_id: Uuid::new_v4().to_string(),
            span_id: Uuid::new_v4().to_string()[..16].to_string(),
            parent_span_id: None,
            attributes: HashMap::new(),
        }
    }

    /// Create a child span
    pub fn child_span(&self) -> Self {
        Self {
            trace_id: self.trace_id.clone(),
            span_id: Uuid::new_v4().to_string()[..16].to_string(),
            parent_span_id: Some(self.span_id.clone()),
            attributes: self.attributes.clone(),
        }
    }

    /// Add an attribute
    pub fn with_attribute(mut self, key: &str, value: &str) -> Self {
        self.attributes.insert(key.to_string(), value.to_string());
        self
    }

    /// Add an attribute
    pub fn set_attribute(&mut self, key: &str, value: &str) {
        self.attributes.insert(key.to_string(), value.to_string());
    }

    /// Headers for propagation
    pub fn to_headers(&self) -> HashMap<String, String> {
        let mut headers = HashMap::new();
        
        // Format W3C Trace Context
        let traceparent = format!(
            "00-{}-{}-01",
            self.trace_id,
            self.span_id
        );
        headers.insert("traceparent".to_string(), traceparent);
        
        // Tracestate (optionnel)
        if !self.attributes.is_empty() {
            let tracestate: String = self.attributes
                .iter()
                .map(|(k, v)| format!("{}={}", k, v))
                .collect::<Vec<_>>()
                .join(",");
            headers.insert("tracestate".to_string(), tracestate);
        }
        
        headers
    }

    /// Parse from headers
    pub fn from_headers(headers: &HashMap<String, String>) -> Option<Self> {
        let traceparent = headers.get("traceparent")?;
        let parts: Vec<&str> = traceparent.split('-').collect();
        
        if parts.len() != 4 || parts[0] != "00" {
            return None;
        }
        
        let trace_id = parts[1].to_string();
        let span_id = parts[2].to_string();
        
        let mut ctx = Self {
            trace_id,
            span_id: Uuid::new_v4().to_string()[..16].to_string(),
            parent_span_id: Some(span_id),
            attributes: HashMap::new(),
        };
        
        // Parse tracestate
        if let Some(tracestate) = headers.get("tracestate") {
            for pair in tracestate.split(',') {
                if let Some((key, value)) = pair.split_once('=') {
                    ctx.attributes.insert(key.trim().to_string(), value.trim().to_string());
                }
            }
        }
        
        Some(ctx)
    }
}

/// Trace span
#[derive(Debug)]
pub struct Span {
    /// Context
    pub context: TraceContext,
    /// Span name
    pub name: String,
    /// Start timestamp
    pub started_at: std::time::Instant,
    /// End timestamp
    pub ended_at: Option<std::time::Instant>,
    /// Status
    pub status: SpanStatus,
    /// Events
    pub events: Vec<SpanEvent>,
}

impl Span {
    /// New span
    pub fn new(context: TraceContext, name: &str) -> Self {
        Self {
            context,
            name: name.to_string(),
            started_at: std::time::Instant::now(),
            ended_at: None,
            status: SpanStatus::Ok,
            events: Vec::new(),
        }
    }

    /// Add an event
    pub fn add_event(&mut self, name: &str, attributes: HashMap<String, String>) {
        self.events.push(SpanEvent {
            name: name.to_string(),
            timestamp: std::time::Instant::now(),
            attributes,
        });
    }

    /// End the span successfully
    pub fn end_ok(&mut self) {
        self.ended_at = Some(std::time::Instant::now());
        self.status = SpanStatus::Ok;
    }

    /// End the span with error
    pub fn end_error(&mut self, message: &str) {
        self.ended_at = Some(std::time::Instant::now());
        self.status = SpanStatus::Error(message.to_string());
    }

    /// Span duration (ms)
    pub fn duration_ms(&self) -> u64 {
        let end = self.ended_at.unwrap_or_else(std::time::Instant::now);
        end.duration_since(self.started_at).as_millis() as u64
    }
}

/// Span status
#[derive(Debug, Clone)]
pub enum SpanStatus {
    /// Successful operation
    Ok,
    /// Failed operation with error message
    Error(String),
}

/// Span event
#[derive(Debug)]
pub struct SpanEvent {
    /// Event name
    pub name: String,
    /// Event timestamp
    pub timestamp: std::time::Instant,
    /// Event attributes
    pub attributes: HashMap<String, String>,
}