use std::sync::Arc;

use axum::{middleware, Extension};
use tokio::sync::RwLock;
use tower_http::cors::CorsLayer;
use tower_http::trace::TraceLayer;

use crate::agent::ComplianceAgent;
use crate::api::auth_middleware::{require_jwt_auth, JwksState};
use crate::api::routes;
use crate::error::AgentError;

pub async fn start_api_server(agent: ComplianceAgent, port: u16) -> Result<(), AgentError> {
    let mut app = routes::build_router()
        .layer(Extension(Arc::new(agent.clone())))
        .layer(CorsLayer::permissive())
        .layer(TraceLayer::new_for_http());

    if let (Some(kc_url), Some(kc_realm)) =
        (&agent.config.keycloak_url, &agent.config.keycloak_realm)
    {
        let jwks_url = format!("{kc_url}/realms/{kc_realm}/protocol/openid-connect/certs");
        let jwks_state = JwksState {
            jwks: Arc::new(RwLock::new(None)),
            jwks_url,
        };
        tracing::info!("Keycloak JWT auth enabled for realm '{kc_realm}'");
        app = app
            .layer(Extension(jwks_state))
            .layer(middleware::from_fn(require_jwt_auth));
    } else {
        tracing::warn!("Keycloak not configured - API endpoints are unprotected");
    }

    let addr = format!("0.0.0.0:{port}");
    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .map_err(|e| AgentError::Other(format!("Failed to bind to {addr}: {e}")))?;

    tracing::info!("REST API listening on {addr}");
    axum::serve(listener, app)
        .await
        .map_err(|e| AgentError::Other(format!("API server error: {e}")))?;

    Ok(())
}