otelite-api 0.1.24

use crate::server::{AppState, QueryCache};
use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    response::IntoResponse,
    Json,
};
use otelite_core::api::{ErrorResponse, SpanEntry, TraceDetail, TraceEntry, TracesResponse};
use otelite_core::query::{Operator, QueryPredicate, QueryValue};
use otelite_core::storage::QueryParams;
use otelite_core::telemetry::Span;
use serde::{Deserialize, Serialize};

/// Query parameters for trace listing
#[derive(Debug, Deserialize, Serialize, utoipa::IntoParams)]
pub struct TracesQuery {
    /// Filter by trace ID
    #[serde(default)]
    pub trace_id: Option<String>,

    /// Filter by service name
    #[serde(default)]
    pub service: Option<String>,

    /// Filter by resource attribute (format: key=value)
    #[serde(default)]
    pub resource: Option<String>,

    /// Full-text search in span names
    #[serde(default)]
    pub search: Option<String>,

    /// Filter by session ID (session.id attribute)
    #[serde(default)]
    pub session_id: Option<String>,

    /// Filter by gen_ai.conversation.id attribute
    #[serde(default)]
    pub conversation_id: Option<String>,

    /// Filter by model (matches gen_ai.request.model)
    #[serde(default)]
    pub model: Option<String>,

    /// JSON-encoded list of attribute filters: [{"key":"x","op":"=","value":"y"}, ...]
    /// Supported ops: "=" (Equal), "!=" (NotEqual)
    #[serde(default)]
    pub attrs: Option<String>,

    /// Start time (Unix timestamp in nanoseconds)
    #[serde(default)]
    pub start_time: Option<i64>,

    /// End time (Unix timestamp in nanoseconds)
    #[serde(default)]
    pub end_time: Option<i64>,

    /// Maximum number of results (default: 100, max: 1000)
    #[serde(default = "default_limit")]
    pub limit: usize,

    /// Offset for pagination
    #[serde(default)]
    pub offset: usize,
}

fn default_limit() -> usize {
    100
}

/// Handler for GET /api/traces
#[utoipa::path(
    get,
    path = "/api/traces",
    params(TracesQuery),
    responses(
        (status = 200, description = "List of traces matching query", body = TracesResponse),
        (status = 500, description = "Internal server error", body = ErrorResponse)
    ),
    tag = "traces"
)]
pub async fn list_traces(
    State(state): State<AppState>,
    Query(params): Query<TracesQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    // Check cache first
    let cache_key = QueryCache::make_key(&params);
    if let Some(cached) = state.cache.traces.get(&cache_key) {
        return Ok((
            StatusCode::OK,
            [("content-type", "application/json")],
            cached,
        )
            .into_response());
    }

    // Validate and cap limit
    let limit = params.limit.min(1000);

    // Build query parameters
    let mut query = QueryParams {
        start_time: params.start_time,
        end_time: params.end_time,
        trace_id: params.trace_id.clone(),
        ..Default::default()
    };

    if let Some(ref sid) = params.session_id {
        if !sid.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "session.id".to_string(),
                operator: Operator::Equal,
                value: QueryValue::String(sid.clone()),
            });
        }
    }

    if let Some(ref cid) = params.conversation_id {
        if !cid.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "gen_ai.conversation.id".to_string(),
                operator: Operator::Equal,
                value: QueryValue::String(cid.clone()),
            });
        }
    }

    if let Some(ref model) = params.model {
        if !model.is_empty() {
            query.predicates.push(QueryPredicate {
                field: "gen_ai.request.model".to_string(),
                operator: Operator::Equal,
                value: QueryValue::String(model.clone()),
            });
        }
    }

    if let Some(ref attrs_json) = params.attrs {
        if !attrs_json.is_empty() {
            #[derive(serde::Deserialize)]
            struct AttrFilter {
                key: String,
                op: String,
                value: Option<String>,
            }
            if let Ok(filters) = serde_json::from_str::<Vec<AttrFilter>>(attrs_json) {
                for f in filters {
                    let op = match f.op.as_str() {
                        "=" => Operator::Equal,
                        "!=" => Operator::NotEqual,
                        _ => continue,
                    };
                    if let Some(v) = f.value {
                        query.predicates.push(QueryPredicate {
                            field: f.key,
                            operator: op,
                            value: QueryValue::String(v),
                        });
                    }
                }
            }
        }
    }

    // Query spans from storage — two-step: get N most-recent trace IDs, then all their spans
    let spans = state
        .storage
        .query_spans_for_trace_list(&query, limit)
        .await
        .map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(ErrorResponse::storage_error(format!("query spans: {}", e))),
            )
        })?;

    // Note: spans do not carry resource attributes in the current data model
    // (resource is on Trace, not Span). The resource query param is accepted
    // for forward compatibility but no filtering is applied here.

    // Group spans by trace_id
    let mut traces_map: std::collections::HashMap<String, Vec<Span>> =
        std::collections::HashMap::new();
    for span in spans {
        traces_map
            .entry(span.trace_id.clone())
            .or_default()
            .push(span);
    }

    // Convert to trace entries
    let mut trace_entries: Vec<TraceEntry> = traces_map
        .into_iter()
        .map(|(trace_id, spans)| {
            let start_time = spans.iter().map(|s| s.start_time).min().unwrap_or(0);
            let end_time = spans.iter().map(|s| s.end_time).max().unwrap_or(0);
            let duration = end_time - start_time;

            let root_span = spans
                .iter()
                .find(|s| s.parent_span_id.is_none())
                .or_else(|| spans.first());

            let root_span_name = root_span
                .map(|s| s.name.clone())
                .unwrap_or_else(|| "Unknown".to_string());

            let service_names: Vec<String> = {
                let mut names: Vec<String> = spans
                    .iter()
                    .filter_map(|s| s.resource.as_ref())
                    .filter_map(|r| r.attributes.get("service.name"))
                    .cloned()
                    .collect::<std::collections::HashSet<_>>()
                    .into_iter()
                    .collect();
                names.sort();
                names
            };

            let has_errors = spans.iter().any(|s| {
                matches!(
                    s.status.code,
                    otelite_core::telemetry::trace::StatusCode::Error
                )
            });

            TraceEntry {
                trace_id,
                root_span_name,
                start_time,
                duration,
                span_count: spans.len(),
                service_names,
                has_errors,
            }
        })
        .collect();

    // Sort by start time (newest first)
    trace_entries.sort_by_key(|b| std::cmp::Reverse(b.start_time));

    // Apply pagination
    let total = trace_entries.len();
    let paginated_traces: Vec<TraceEntry> = trace_entries
        .into_iter()
        .skip(params.offset)
        .take(limit)
        .collect();

    let response = TracesResponse {
        traces: paginated_traces,
        total,
        limit,
        offset: params.offset,
    };

    // Cache the response
    if let Ok(json) = serde_json::to_string(&response) {
        state.cache.traces.insert(cache_key, json.clone());
        Ok((StatusCode::OK, [("content-type", "application/json")], json).into_response())
    } else {
        Ok(Json(response).into_response())
    }
}

/// Handler for GET /api/traces/:trace_id
#[utoipa::path(
    get,
    path = "/api/traces/{trace_id}",
    params(
        ("trace_id" = String, Path, description = "Trace ID")
    ),
    responses(
        (status = 200, description = "Trace details with all spans", body = TraceDetail),
        (status = 404, description = "Trace not found", body = ErrorResponse),
        (status = 500, description = "Internal server error", body = ErrorResponse)
    ),
    tag = "traces"
)]
pub async fn get_trace(
    State(state): State<AppState>,
    Path(trace_id): Path<String>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    // Query all spans for this trace
    let query = QueryParams {
        trace_id: Some(trace_id.clone()),
        limit: Some(1000), // Max spans per trace
        ..Default::default()
    };

    let spans = state.storage.query_spans(&query).await.map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse::storage_error(format!(
                "query trace spans: {}",
                e
            ))),
        )
    })?;

    if spans.is_empty() {
        return Err((
            StatusCode::NOT_FOUND,
            Json(ErrorResponse::not_found(format!("Trace '{}'", trace_id))),
        ));
    }

    let start_time = spans.iter().map(|s| s.start_time).min().unwrap_or(0);
    let end_time = spans.iter().map(|s| s.end_time).max().unwrap_or(0);
    let duration = end_time - start_time;

    let service_names: Vec<String> = {
        let mut names: Vec<String> = spans
            .iter()
            .filter_map(|s| s.resource.as_ref())
            .filter_map(|r| r.attributes.get("service.name"))
            .cloned()
            .collect::<std::collections::HashSet<_>>()
            .into_iter()
            .collect();
        names.sort();
        names
    };

    let span_entries: Vec<SpanEntry> = spans.into_iter().map(SpanEntry::from).collect();

    let span_count = span_entries.len();

    let trace_detail = TraceDetail {
        trace_id,
        spans: span_entries,
        start_time,
        end_time,
        duration,
        span_count,
        service_names,
    };

    Ok(Json(trace_detail))
}

/// Export format for traces
#[derive(Debug, Deserialize, utoipa::IntoParams)]
pub struct ExportQuery {
    /// Export format: "json"
    #[serde(default = "default_format")]
    pub format: String,

    /// Same filters as TracesQuery
    #[serde(flatten)]
    pub filters: TracesQuery,
}

fn default_format() -> String {
    "json".to_string()
}

/// Handler for GET /api/traces/export
#[utoipa::path(
    get,
    path = "/api/traces/export",
    params(ExportQuery),
    responses(
        (status = 200, description = "Exported traces in JSON format"),
        (status = 400, description = "Invalid format parameter", body = ErrorResponse),
        (status = 500, description = "Internal server error", body = ErrorResponse)
    ),
    tag = "traces"
)]
pub async fn export_traces(
    State(state): State<AppState>,
    Query(params): Query<ExportQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    // Build query parameters (no limit for export, but cap at 10000)
    let query = QueryParams {
        start_time: params.filters.start_time,
        end_time: params.filters.end_time,
        limit: Some(10000),
        trace_id: params.filters.trace_id.clone().filter(|s| !s.is_empty()),
        ..Default::default()
    };

    let spans = state.storage.query_spans(&query).await.map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse::storage_error(format!(
                "export traces: {}",
                e
            ))),
        )
    })?;

    match params.format.as_str() {
        "json" => {
            let span_entries: Vec<SpanEntry> = spans.into_iter().map(SpanEntry::from).collect();

            Ok((
                [
                    ("Content-Type", "application/json"),
                    (
                        "Content-Disposition",
                        "attachment; filename=\"traces.json\"",
                    ),
                ],
                Json(span_entries),
            )
                .into_response())
        },
        _ => Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse::bad_request(
                "Invalid format parameter. Use 'json'",
            )),
        )),
    }
}