src/main.rs

use axum::{routing::{get, post}, Router, response::Json};
use serde::Serialize;
use tower_http::trace::TraceLayer;
use tracing_subscriber::EnvFilter;

mod handlers;
mod inference;
mod shield;
mod web3;
mod orchestrator;
mod proof;
mod api;
mod federation;

// use the handler directly from the crate root, not from a separate library
use handlers::triage;

#[derive(Serialize)]
struct StatusResponse {
    status: String,
    model: String,
    device: String,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    dotenvy::dotenv().ok();
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env().add_directive("rustvital_amd=debug".parse()?))
        .init();

    // Init zk-lite signing key
    proof::init_signing_key("0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef");

    let port = std::env::var("PORT").unwrap_or_else(|_| "3000".to_string());
    let addr = format!("0.0.0.0:{}", port);
    tracing::info!("Starting RustVital-AMD server on {}", addr);

    let app = Router::new()
        .route("/", get(serve_ui))
        .route("/health", get(|| async { "healthy" }))
        .route("/status", get(status))
        .route("/triage", post(triage::handle))
        .route("/triage/stream", get(api::stream::triage_stream))
        .route("/trigger-federated-tune", post(federation::trigger_tune))
        .route("/federation/round", get(federation::latest_round))
        .layer(TraceLayer::new_for_http());

    let listener = tokio::net::TcpListener::bind(&addr).await?;
    axum::serve(listener, app).await?;
    Ok(())
}

async fn serve_ui() -> axum::response::Html<&'static str> {
    axum::response::Html(include_str!("../static/index.html"))
}

async fn status() -> Json<StatusResponse> {
    let device = if std::env::var("ENABLE_ROCM").unwrap_or_default() == "1" {
        "ROCm/HIP (MI300X)"
    } else if std::env::var("VLLM_URL").is_ok() {
        "ROCm/MI300X (vLLM connected)"
    } else {
        "CPU (fallback)"
    };
    let model = std::env::var("FORCE_0_5B").map_or("7B (Qwen2.5-7B-Instruct)".to_string(), |_| "0.5B (Qwen2.5-0.5B-Instruct)".to_string());
    Json(StatusResponse {
        status: "running".to_string(),
        model,
        device: device.to_string(),
    })
}