use std::sync::Arc;

use axum::Router;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use sea_orm::{ConnectionTrait, DbBackend, Statement};

use crate::{AppDatabase, AppRedis, AppState};

#[derive(Clone)]
struct HealthState {
    database: Option<Arc<AppDatabase>>,
    redis: AppRedis,
    kenauth: Option<Arc<kenauth::KenAuthClient>>,
}

async fn healthz() -> Response {
    (StatusCode::OK, "ok").into_response()
}

async fn readyz(State(state): State<HealthState>) -> Response {
    let db_ok = match state.database.as_ref() {
        None => true,
        Some(db) => db
            .read()
            .query_one_raw(Statement::from_sql_and_values(
                DbBackend::Postgres,
                "SELECT 1",
                [],
            ))
            .await
            .is_ok(),
    };
    let mut conn = state.redis.connection_clone();
    let redis_ok: bool = match &mut conn {
        crate::redis::RedisConnection::Standalone(c) => {
            redis::cmd("PING").query_async::<String>(c).await.is_ok()
        }
        crate::redis::RedisConnection::Cluster(c) => {
            redis::cmd("PING").query_async::<String>(c).await.is_ok()
        }
    };
    let kenauth_ok = match state.kenauth.as_ref() {
        None => true,
        Some(client) => client.check_health().await.is_ok(),
    };
    if db_ok && redis_ok && kenauth_ok {
        (StatusCode::OK, "ready").into_response()
    } else {
        (StatusCode::SERVICE_UNAVAILABLE, "not ready").into_response()
    }
}

pub fn router_with(
    database: Option<Arc<AppDatabase>>,
    redis: AppRedis,
    kenauth: Option<Arc<kenauth::KenAuthClient>>,
) -> Router {
    Router::new()
        .route("/healthz", axum::routing::get(healthz))
        .route("/readyz", axum::routing::get(readyz))
        .with_state(HealthState {
            database,
            redis,
            kenauth,
        })
}

pub fn router(state: AppState) -> Router {
    router_with(
        Some(state.database.clone()),
        state.redis.clone(),
        state.kenauth.clone(),
    )
}