bloop-sdk 0.2.0

use std::sync::Mutex;
use crate::types::{TraceData, TraceStatus, SpanData};

/// A live LLM trace that collects spans.
///
/// Traces use interior mutability via Mutex so that spans can add themselves
/// to the trace without requiring exclusive (&mut) access. This enables the
/// ergonomic pattern:
///
/// ```rust,ignore
/// let trace = client.start_trace("chat");
/// let mut span = Span::new(&trace, SpanType::Generation, "llm");
/// span.end();  // adds to trace internally
/// trace.end();
/// ```
pub struct Trace {
    id: String,
    name: String,
    status: Mutex<TraceStatus>,
    started_at: i64,
    ended_at: Mutex<Option<i64>>,
    session_id: Option<String>,
    user_id: Option<String>,
    input: Option<String>,
    output: Mutex<Option<String>>,
    metadata: Option<serde_json::Value>,
    prompt_name: Option<String>,
    prompt_version: Option<String>,
    spans: Mutex<Vec<SpanData>>,
}

impl Trace {
    /// Create a new trace with the given name.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string().replace("-", ""),
            name: name.into(),
            status: Mutex::new(TraceStatus::Running),
            started_at: now_millis(),
            ended_at: Mutex::new(None),
            session_id: None,
            user_id: None,
            input: None,
            output: Mutex::new(None),
            metadata: None,
            prompt_name: None,
            prompt_version: None,
            spans: Mutex::new(Vec::new()),
        }
    }

    // ── Builder-pattern setters (consume self) ──

    pub fn session_id(mut self, id: impl Into<String>) -> Self {
        self.session_id = Some(id.into());
        self
    }

    pub fn user_id(mut self, id: impl Into<String>) -> Self {
        self.user_id = Some(id.into());
        self
    }

    pub fn input_text(mut self, input: impl Into<String>) -> Self {
        self.input = Some(input.into());
        self
    }

    pub fn prompt_name(mut self, name: impl Into<String>) -> Self {
        self.prompt_name = Some(name.into());
        self
    }

    pub fn prompt_version(mut self, version: impl Into<String>) -> Self {
        self.prompt_version = Some(version.into());
        self
    }

    // ── Mutating methods ──

    pub fn set_output(&self, output: impl Into<String>) {
        *self.output.lock().unwrap() = Some(output.into());
    }

    pub fn set_status(&self, status: TraceStatus) {
        *self.status.lock().unwrap() = status;
    }

    pub fn end(&self) {
        let mut ended = self.ended_at.lock().unwrap();
        if ended.is_none() {
            *ended = Some(now_millis());
        }
        let mut status = self.status.lock().unwrap();
        if *status == TraceStatus::Running {
            *status = TraceStatus::Completed;
        }
    }

    // ── Internal: called by Span::end() ──

    pub(crate) fn add_span(&self, span_data: SpanData) {
        self.spans.lock().unwrap().push(span_data);
    }

    // ── Accessors ──

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn started_at(&self) -> i64 {
        self.started_at
    }

    pub fn span_count(&self) -> usize {
        self.spans.lock().unwrap().len()
    }

    pub fn spans(&self) -> Vec<SpanData> {
        self.spans.lock().unwrap().clone()
    }

    /// Convert to owned TraceData for serialization/sending.
    pub fn to_data(&self) -> TraceData {
        TraceData {
            id: self.id.clone(),
            name: self.name.clone(),
            status: *self.status.lock().unwrap(),
            started_at: self.started_at,
            ended_at: *self.ended_at.lock().unwrap(),
            spans: self.spans.lock().unwrap().clone(),
            session_id: self.session_id.clone(),
            user_id: self.user_id.clone(),
            input: self.input.clone(),
            output: self.output.lock().unwrap().clone(),
            metadata: self.metadata.clone(),
            prompt_name: self.prompt_name.clone(),
            prompt_version: self.prompt_version.clone(),
        }
    }
}

impl std::fmt::Debug for Trace {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Trace")
            .field("id", &self.id)
            .field("name", &self.name)
            .field("started_at", &self.started_at)
            .field("span_count", &self.span_count())
            .finish()
    }
}

fn now_millis() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64
}