revoke-trace 0.3.0

use opentelemetry::trace::{SpanKind as OtelSpanKind, Status};
use std::collections::HashMap;
use std::time::SystemTime;
use tracing::Level;

/// Span 类型
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpanKind {
    /// 服务端处理请求
    Server,
    /// 客户端发送请求
    Client,
    /// 生产者发送消息
    Producer,
    /// 消费者接收消息
    Consumer,
    /// 内部操作
    Internal,
}

impl From<SpanKind> for OtelSpanKind {
    fn from(kind: SpanKind) -> Self {
        match kind {
            SpanKind::Server => OtelSpanKind::Server,
            SpanKind::Client => OtelSpanKind::Client,
            SpanKind::Producer => OtelSpanKind::Producer,
            SpanKind::Consumer => OtelSpanKind::Consumer,
            SpanKind::Internal => OtelSpanKind::Internal,
        }
    }
}

/// Span 状态
#[derive(Debug, Clone)]
pub struct SpanStatus {
    code: SpanStatusCode,
    message: String,
}

/// Span 状态码
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpanStatusCode {
    /// 成功
    Ok,
    /// 错误
    Error,
    /// 未设置
    Unset,
}

impl SpanStatus {
    /// 创建成功状态
    pub fn ok() -> Self {
        Self {
            code: SpanStatusCode::Ok,
            message: String::new(),
        }
    }

    /// 创建错误状态
    pub fn error(message: impl Into<String>) -> Self {
        Self {
            code: SpanStatusCode::Error,
            message: message.into(),
        }
    }

    /// 创建未设置状态
    pub fn unset() -> Self {
        Self {
            code: SpanStatusCode::Unset,
            message: String::new(),
        }
    }
}

impl From<SpanStatus> for Status {
    fn from(status: SpanStatus) -> Self {
        match status.code {
            SpanStatusCode::Ok => Status::Ok,
            SpanStatusCode::Error => Status::error(status.message),
            SpanStatusCode::Unset => Status::Unset,
        }
    }
}

/// Span 构建器
pub struct SpanBuilder {
    name: String,
    kind: SpanKind,
    attributes: HashMap<String, String>,
    events: Vec<SpanEvent>,
    links: Vec<SpanLink>,
    start_time: Option<SystemTime>,
    level: Level,
}

impl SpanBuilder {
    /// 创建新的 Span 构建器
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            kind: SpanKind::Internal,
            attributes: HashMap::new(),
            events: Vec::new(),
            links: Vec::new(),
            start_time: None,
            level: Level::INFO,
        }
    }

    /// 设置 Span 类型
    pub fn with_kind(mut self, kind: SpanKind) -> Self {
        self.kind = kind;
        self
    }

    /// 添加属性
    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.attributes.insert(key.into(), value.into());
        self
    }

    /// 批量添加属性
    pub fn with_attributes<I, K, V>(mut self, attributes: I) -> Self
    where
        I: IntoIterator<Item = (K, V)>,
        K: Into<String>,
        V: Into<String>,
    {
        for (key, value) in attributes {
            self.attributes.insert(key.into(), value.into());
        }
        self
    }

    /// 添加事件
    pub fn with_event(mut self, event: SpanEvent) -> Self {
        self.events.push(event);
        self
    }

    /// 添加链接
    pub fn with_link(mut self, link: SpanLink) -> Self {
        self.links.push(link);
        self
    }

    /// 设置开始时间
    pub fn with_start_time(mut self, time: SystemTime) -> Self {
        self.start_time = Some(time);
        self
    }

    /// 设置日志级别
    pub fn with_level(mut self, level: Level) -> Self {
        self.level = level;
        self
    }

    /// 构建并开始 Span
    pub fn start(self) -> tracing::Span {
        match self.level {
            Level::TRACE => {
                tracing::trace_span!(
                    target: "revoke_trace",
                    "{}",
                    self.name,
                    otel.kind = ?self.kind
                )
            }
            Level::DEBUG => {
                tracing::debug_span!(
                    target: "revoke_trace",
                    "{}",
                    self.name,
                    otel.kind = ?self.kind
                )
            }
            Level::INFO => {
                tracing::info_span!(
                    target: "revoke_trace",
                    "{}",
                    self.name,
                    otel.kind = ?self.kind
                )
            }
            Level::WARN => {
                tracing::warn_span!(
                    target: "revoke_trace",
                    "{}",
                    self.name,
                    otel.kind = ?self.kind
                )
            }
            Level::ERROR => {
                tracing::error_span!(
                    target: "revoke_trace",
                    "{}",
                    self.name,
                    otel.kind = ?self.kind
                )
            }
        }
    }
}

/// Span 事件
#[derive(Debug, Clone)]
pub struct SpanEvent {
    name: String,
    timestamp: SystemTime,
    attributes: HashMap<String, String>,
}

impl SpanEvent {
    /// 创建新的事件
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            timestamp: SystemTime::now(),
            attributes: HashMap::new(),
        }
    }

    /// 添加属性
    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.attributes.insert(key.into(), value.into());
        self
    }

    /// 设置时间戳
    pub fn with_timestamp(mut self, timestamp: SystemTime) -> Self {
        self.timestamp = timestamp;
        self
    }
}

/// Span 链接
#[derive(Debug, Clone)]
pub struct SpanLink {
    #[allow(dead_code)]
    context: crate::context::TraceContext,
    attributes: HashMap<String, String>,
}

impl SpanLink {
    /// 创建新的链接
    pub fn new(context: crate::context::TraceContext) -> Self {
        Self {
            context,
            attributes: HashMap::new(),
        }
    }

    /// 添加属性
    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.attributes.insert(key.into(), value.into());
        self
    }
}

/// Span 扩展功能
pub trait SpanExt {
    /// 记录异常
    fn record_exception(&self, error: &dyn std::error::Error);

    /// 设置状态
    fn set_status(&self, status: SpanStatus);

    /// 添加事件
    fn add_event(&self, event: SpanEvent);
}

impl SpanExt for tracing::Span {
    fn record_exception(&self, error: &dyn std::error::Error) {
        self.record("exception.type", &error.to_string());
        self.record("exception.message", &error.to_string());

        if let Some(source) = error.source() {
            self.record("exception.stacktrace", &source.to_string());
        }
    }

    fn set_status(&self, status: SpanStatus) {
        self.record("otel.status_code", format!("{:?}", status.code).as_str());
        if !status.message.is_empty() {
            self.record("otel.status_message", &status.message);
        }
    }

    fn add_event(&self, event: SpanEvent) {
        tracing::event!(
            target: "revoke_trace",
            parent: self,
            Level::INFO,
            name = %event.name,
            timestamp = ?event.timestamp,
            ?event.attributes,
            "span event"
        );
    }
}

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

    #[test]
    fn test_span_builder() {
        let _span = SpanBuilder::new("test_operation")
            .with_kind(SpanKind::Server)
            .with_attribute("http.method", "GET")
            .with_attribute("http.url", "/api/users")
            .with_level(Level::DEBUG)
            .start();

        // 如果没有 panic，说明成功创建
    }

    #[test]
    fn test_span_status() {
        let ok_status = SpanStatus::ok();
        let error_status = SpanStatus::error("Something went wrong");

        // 只验证转换不会 panic
        let _otel_status: Status = ok_status.into();
        let _otel_error: Status = error_status.into();
    }
}