arcly-http 0.1.0

//! Health-check registry and `/healthz` / `/readyz` route handlers.
//!
//! Register checks from any plugin or application code via `global().register(...)`.
//! The plugin calls `add_get("/healthz", ...)` and `add_get("/readyz", ...)` to
//! expose them, both backed by `global().run_all()`.

use std::collections::BTreeMap;
use std::sync::{Arc, OnceLock};
use std::time::Instant;

use axum::response::Response;
use dashmap::DashMap;
use futures::future::BoxFuture;
use serde::Serialize;

// ─── Public surface ──────────────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum HealthStatus {
    Healthy,
    Degraded(String),
    Unhealthy(String),
}

/// One named sub-system health probe.
pub trait HealthCheck: Send + Sync + 'static {
    fn check(&self) -> BoxFuture<'_, HealthStatus>;
}

/// Process-wide health-check registry. Accessible via `global()`.
pub struct HealthRegistry {
    checks: DashMap<&'static str, Arc<dyn HealthCheck>>,
    started_at: Instant,
}

impl HealthRegistry {
    fn new() -> Self {
        Self {
            checks: DashMap::new(),
            started_at: Instant::now(),
        }
    }

    /// Register a named health check. Idempotent — re-registration replaces the
    /// previous check for that name.
    pub fn register(&self, name: &'static str, check: impl HealthCheck) {
        self.checks.insert(name, Arc::new(check));
    }

    /// Run all registered checks concurrently and return a map of their results.
    pub async fn run_all(&self) -> BTreeMap<&'static str, HealthStatus> {
        let mut results = BTreeMap::new();
        // Collect futures first so we don't hold the DashMap ref across await.
        let futs: Vec<(&'static str, Arc<dyn HealthCheck>)> = self
            .checks
            .iter()
            .map(|e| (*e.key(), Arc::clone(e.value())))
            .collect();

        for (name, check) in futs {
            results.insert(name, check.check().await);
        }
        results
    }

    pub fn uptime_secs(&self) -> u64 {
        self.started_at.elapsed().as_secs()
    }
}

/// The one process-wide health registry. Initialized on first access.
pub fn global() -> &'static HealthRegistry {
    static REGISTRY: OnceLock<HealthRegistry> = OnceLock::new();
    REGISTRY.get_or_init(HealthRegistry::new)
}

// ─── Route handlers ──────────────────────────────────────────────────────────

#[derive(Serialize)]
struct HealthResponse<'a> {
    status: &'a str,
    checks: BTreeMap<&'static str, serde_json::Value>,
    uptime_secs: u64,
}

async fn run_health_response() -> Response {
    let registry = global();
    let results = registry.run_all().await;

    let mut any_unhealthy = false;
    let mut any_degraded = false;
    let mut checks: BTreeMap<&'static str, serde_json::Value> = BTreeMap::new();
    for (name, status) in &results {
        match status {
            HealthStatus::Unhealthy(_) => any_unhealthy = true,
            HealthStatus::Degraded(_) => any_degraded = true,
            HealthStatus::Healthy => {}
        }
        checks.insert(name, serde_json::to_value(status).unwrap_or_default());
    }

    let overall = if any_unhealthy {
        "unhealthy"
    } else if any_degraded {
        "degraded"
    } else {
        "healthy"
    };
    let http_status = if any_unhealthy { 503u16 } else { 200 };

    let body = serde_json::to_vec(&HealthResponse {
        status: overall,
        checks,
        uptime_secs: registry.uptime_secs(),
    })
    .unwrap_or_default();

    Response::builder()
        .status(http_status)
        .header("Content-Type", "application/json")
        .body(axum::body::Body::from(body))
        .unwrap()
}

fn probe_handler() -> impl Fn(crate::web::context::RequestContext) -> BoxFuture<'static, Response>
       + Send
       + Sync
       + Clone
       + 'static {
    |_ctx| Box::pin(run_health_response())
}

/// Handler for `GET /healthz` — liveness probe.
pub fn healthz_handler(
) -> impl Fn(crate::web::context::RequestContext) -> BoxFuture<'static, Response>
       + Send
       + Sync
       + Clone
       + 'static {
    probe_handler()
}

/// Handler for `GET /readyz` — readiness probe.
pub fn readyz_handler(
) -> impl Fn(crate::web::context::RequestContext) -> BoxFuture<'static, Response>
       + Send
       + Sync
       + Clone
       + 'static {
    probe_handler()
}