//! HTTP `/health` endpoint per `containerization.md §7`.

use std::net::SocketAddr;
use std::sync::Arc;

use anyhow::{Context, Result};
use axum::{routing::get, Json, Router};
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use tracing::{error, info};

use shared::health::AggregatedHealth;

use crate::runtime::Runtime;

#[derive(Clone)]
struct AppState {
    runtime: Arc<Runtime>,
}

pub struct HealthServerHandle {
    shutdown_tx: Option<oneshot::Sender<()>>,
    task: JoinHandle<()>,
}

impl HealthServerHandle {
    pub async fn shutdown(mut self) {
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(());
        }
        if let Err(e) = self.task.await {
            error!(error = ?e, "health server task did not shut down cleanly");
        }
    }
}

/// Start the HTTP health server. Returns once the listener is bound.
pub async fn spawn(bind: String, runtime: Arc<Runtime>) -> Result<HealthServerHandle> {
    let addr: SocketAddr = bind
        .parse()
        .with_context(|| format!("invalid health.bind address: {bind}"))?;

    let state = AppState { runtime };
    let app = Router::new()
        .route("/health", get(health_handler))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind(addr)
        .await
        .with_context(|| format!("binding health server to {addr}"))?;

    info!(%addr, "health server listening");

    let (shutdown_tx, shutdown_rx) = oneshot::channel();
    let task = tokio::spawn(async move {
        if let Err(e) = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = shutdown_rx.await;
            })
            .await
        {
            error!(error = %e, "health server exited with error");
        }
    });

    Ok(HealthServerHandle {
        shutdown_tx: Some(shutdown_tx),
        task,
    })
}

async fn health_handler(
    axum::extract::State(state): axum::extract::State<AppState>,
) -> Json<AggregatedHealth> {
    Json(state.runtime.health_snapshot())
}
