revoke-trace 0.3.0

use axum::{
    Router,
    extract::State,
    http::{HeaderMap, StatusCode},
    response::Json,
    routing::{get, post},
};
use revoke_trace::{
    init_tracer,
    propagator::{HttpHeaders, TracePropagator},
    shutdown_tracer,
    span::{SpanBuilder, SpanKind},
    tracer::TracerConfig,
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tracing::{Instrument, info, instrument};

/// 共享应用状态
#[derive(Clone)]
struct AppState {
    client: reqwest::Client,
    propagator: Arc<TracePropagator>,
}

/// API Gateway 服务
pub async fn run_gateway() -> Result<(), Box<dyn std::error::Error>> {
    let config = TracerConfig {
        service_name: "api-gateway".to_string(),
        service_version: "1.0.0".to_string(),
        environment: "production".to_string(),
        ..Default::default()
    };
    init_tracer(config).await?;

    let state = AppState {
        client: reqwest::Client::new(),
        propagator: Arc::new(TracePropagator::w3c()),
    };

    let app = Router::new()
        .route("/api/products/:id", get(get_product))
        .route("/api/orders", post(create_order))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
    info!("API Gateway listening on :8080");

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;

    shutdown_tracer().await?;
    Ok(())
}

/// 产品服务
pub async fn run_product_service() -> Result<(), Box<dyn std::error::Error>> {
    let config = TracerConfig {
        service_name: "product-service".to_string(),
        service_version: "1.0.0".to_string(),
        environment: "production".to_string(),
        ..Default::default()
    };
    init_tracer(config).await?;

    let app = Router::new()
        .route("/products/:id", get(get_product_handler))
        .route("/products/:id/inventory", get(check_inventory));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:8081").await?;
    info!("Product Service listening on :8081");

    axum::serve(listener, app).await?;
    shutdown_tracer().await?;
    Ok(())
}

/// 库存服务
pub async fn run_inventory_service() -> Result<(), Box<dyn std::error::Error>> {
    let config = TracerConfig {
        service_name: "inventory-service".to_string(),
        service_version: "1.0.0".to_string(),
        environment: "production".to_string(),
        ..Default::default()
    };
    init_tracer(config).await?;

    let app = Router::new().route("/inventory/:product_id", get(get_inventory_handler));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:8082").await?;
    info!("Inventory Service listening on :8082");

    axum::serve(listener, app).await?;
    shutdown_tracer().await?;
    Ok(())
}

// Gateway 处理函数
#[instrument(skip(state))]
async fn get_product(
    State(state): State<AppState>,
    axum::extract::Path(id): axum::extract::Path<u64>,
) -> Result<Json<ProductWithInventory>, StatusCode> {
    info!(product_id = %id, "Fetching product from gateway");

    // 调用产品服务
    let product = call_product_service(&state, id).await?;

    Ok(Json(product))
}

#[instrument(skip(state))]
async fn create_order(
    State(state): State<AppState>,
    Json(order): Json<CreateOrderRequest>,
) -> Result<Json<OrderResponse>, StatusCode> {
    info!(user_id = %order.user_id, "Creating new order");

    // 验证所有产品的库存
    for item in &order.items {
        let product = call_product_service(&state, item.product_id).await?;
        if product.inventory.available < item.quantity {
            return Err(StatusCode::BAD_REQUEST);
        }
    }

    // 创建订单
    let order_id = rand::random::<u64>();
    info!(order_id = %order_id, "Order created successfully");

    Ok(Json(OrderResponse {
        order_id,
        status: "created".to_string(),
        total: order.items.iter().map(|i| i.quantity as f64 * 10.0).sum(),
    }))
}

// 调用产品服务
async fn call_product_service(
    state: &AppState,
    product_id: u64,
) -> Result<ProductWithInventory, StatusCode> {
    let span = SpanBuilder::new("http.request")
        .with_kind(SpanKind::Client)
        .with_attribute("http.method", "GET")
        .with_attribute(
            "http.url",
            &format!("http://localhost:8081/products/{}", product_id),
        )
        .with_attribute("peer.service", "product-service")
        .start();

    async move {
        // 准备请求头，注入追踪上下文
        let mut headers = HeaderMap::new();
        let mut http_headers = HttpHeaders::new();

        // 获取当前追踪上下文并注入到请求头
        let current_context = tracing::Span::current();
        state
            .propagator
            .inject(&current_context.context(), &mut http_headers);

        // 转换为 HeaderMap
        for (key, value) in http_headers {
            if let (Ok(name), Ok(value)) = (
                axum::http::header::HeaderName::from_bytes(key.as_bytes()),
                axum::http::header::HeaderValue::from_str(&value),
            ) {
                headers.insert(name, value);
            }
        }

        // 发送请求
        let response = state
            .client
            .get(format!("http://localhost:8081/products/{}", product_id))
            .headers(headers)
            .send()
            .await
            .map_err(|_| StatusCode::BAD_GATEWAY)?;

        if response.status().is_success() {
            response
                .json::<ProductWithInventory>()
                .await
                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
        } else {
            Err(StatusCode::from_u16(response.status().as_u16())
                .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR))
        }
    }
    .instrument(span)
    .await
}

// 产品服务处理函数
#[instrument]
async fn get_product_handler(
    axum::extract::Path(id): axum::extract::Path<u64>,
) -> Result<Json<ProductWithInventory>, StatusCode> {
    info!(product_id = %id, "Fetching product details");

    // 获取产品信息
    let product = Product {
        id,
        name: format!("Product {}", id),
        price: 99.99,
        category: "Electronics".to_string(),
    };

    // 调用库存服务
    let inventory = call_inventory_service(id).await?;

    Ok(Json(ProductWithInventory { product, inventory }))
}

#[instrument]
async fn check_inventory(
    axum::extract::Path(id): axum::extract::Path<u64>,
) -> Result<Json<Inventory>, StatusCode> {
    call_inventory_service(id).await.map(Json)
}

// 调用库存服务
async fn call_inventory_service(product_id: u64) -> Result<Inventory, StatusCode> {
    let span = SpanBuilder::new("http.request")
        .with_kind(SpanKind::Client)
        .with_attribute("http.method", "GET")
        .with_attribute(
            "http.url",
            &format!("http://localhost:8082/inventory/{}", product_id),
        )
        .with_attribute("peer.service", "inventory-service")
        .start();

    async move {
        // 这里简化了，实际应该也要传播追踪上下文
        let client = reqwest::Client::new();
        let response = client
            .get(format!("http://localhost:8082/inventory/{}", product_id))
            .send()
            .await
            .map_err(|_| StatusCode::BAD_GATEWAY)?;

        if response.status().is_success() {
            response
                .json::<Inventory>()
                .await
                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
        } else {
            Err(StatusCode::NOT_FOUND)
        }
    }
    .instrument(span)
    .await
}

// 库存服务处理函数
#[instrument]
async fn get_inventory_handler(
    axum::extract::Path(product_id): axum::extract::Path<u64>,
) -> Json<Inventory> {
    info!(product_id = %product_id, "Checking inventory");

    // 模拟库存查询
    let inventory = Inventory {
        product_id,
        available: rand::random::<u32>() % 100,
        reserved: rand::random::<u32>() % 20,
        warehouse_location: "Warehouse A".to_string(),
    };

    Json(inventory)
}

// 数据模型
#[derive(Serialize, Deserialize)]
struct Product {
    id: u64,
    name: String,
    price: f64,
    category: String,
}

#[derive(Serialize, Deserialize)]
struct Inventory {
    product_id: u64,
    available: u32,
    reserved: u32,
    warehouse_location: String,
}

#[derive(Serialize, Deserialize)]
struct ProductWithInventory {
    product: Product,
    inventory: Inventory,
}

#[derive(Deserialize)]
struct CreateOrderRequest {
    user_id: u64,
    items: Vec<OrderItem>,
}

#[derive(Deserialize)]
struct OrderItem {
    product_id: u64,
    quantity: u32,
}

#[derive(Serialize)]
struct OrderResponse {
    order_id: u64,
    status: String,
    total: f64,
}

async fn shutdown_signal() {
    tokio::signal::ctrl_c()
        .await
        .expect("failed to install Ctrl+C handler");
}

// 主函数示例 - 启动所有服务
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 在实际部署中，这些服务会在不同的进程/容器中运行
    // 这里为了演示，我们在同一个进程中启动它们

    let gateway = tokio::spawn(run_gateway());
    let product = tokio::spawn(run_product_service());
    let inventory = tokio::spawn(run_inventory_service());

    // 等待任一服务退出
    tokio::select! {
        _ = gateway => {},
        _ = product => {},
        _ = inventory => {},
    }

    Ok(())
}