use axum::{extract::Request, http::StatusCode, middleware::Next, response::Response};

use crate::auth::{
    jwt::validate_token,
    keycloak::{get_jwks, refresh_jwks},
};

pub async fn auth_middleware(mut request: Request, next: Next) -> Result<Response, StatusCode> {
    dbg!("Middleware hit");

    let headers = request.headers();

    dbg!("Headers extracted");

    let auth_header = headers.get("authorization").and_then(|v| v.to_str().ok());

    dbg!("Auth header: {:?}", auth_header);

    let auth_header = auth_header.ok_or(StatusCode::UNAUTHORIZED)?;

    let token = auth_header
        .strip_prefix("Bearer ")
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let jwks = get_jwks().await.map_err(|_| StatusCode::UNAUTHORIZED)?;

    match validate_token(token, &jwks) {
        Ok(claims) => {
            dbg!("Token valid");

            request.extensions_mut().insert(claims);

            Ok(next.run(request).await)
        }
        Err(_) => {
            let jwks = refresh_jwks().await.map_err(|_| StatusCode::UNAUTHORIZED)?;

            match validate_token(token, &jwks) {
                Ok(claims) => {
                    request.extensions_mut().insert(claims);
                    Ok(next.run(request).await)
                }
                Err(_) => Err(StatusCode::UNAUTHORIZED),
            }
        }
    }
}