use crate::config::Config;
use crate::crypto::CryptoProvider;
use crate::db::DbPool;
use crate::errors::CognitoError;
use crate::operations;
use axum::body::Bytes;
use axum::extract::State;
use axum::http::HeaderMap;
use axum::response::Response;
use axum::routing::{get, post};
use axum::Router;
use std::sync::Arc;

#[derive(Clone)]
pub struct AppState {
    pub db: DbPool,
    pub crypto: CryptoProvider,
    pub config: Config,
}

pub fn build_router(state: AppState) -> Router {
    Router::new()
        .route("/", post(dispatch))
        .route("/.well-known/jwks.json", get(jwks))
        .route("/health", get(health))
        .with_state(Arc::new(state))
}

async fn health() -> &'static str {
    "ok"
}

async fn jwks(State(state): State<Arc<AppState>>) -> axum::Json<serde_json::Value> {
    axum::Json(state.crypto.jwks().clone())
}

async fn dispatch(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    body: Bytes,
) -> Result<Response, CognitoError> {
    let target = headers
        .get("x-amz-target")
        .and_then(|v| v.to_str().ok())
        .ok_or_else(|| {
            CognitoError::InvalidParameter("Missing X-Amz-Target header".to_string())
        })?;

    let operation = target
        .rsplit('.')
        .next()
        .ok_or_else(|| CognitoError::InvalidParameter("Invalid X-Amz-Target".to_string()))?;

    tracing::info!(operation = operation, "handling request");

    let json_body: serde_json::Value = if body.is_empty() {
        serde_json::Value::Object(serde_json::Map::new())
    } else {
        serde_json::from_slice(&body).map_err(|e| {
            CognitoError::InvalidParameter(format!("Invalid JSON: {e}"))
        })?
    };

    operations::dispatch(operation, json_body, &state).await
}
