use auth_kit::jwt::{Claims, JwtVerifyConfig};
use axum::response::IntoResponse;
use axum::{Json, Router, routing::get};
use base64::Engine as _;
use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
use rsa::pkcs1::{EncodeRsaPrivateKey, EncodeRsaPublicKey};
use rsa::rand_core::OsRng;
use rsa::traits::PublicKeyParts;
use rsa::{RsaPrivateKey, pkcs1::LineEnding};
use serde::Serialize;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

#[derive(Serialize, Clone)]
struct Jwks {
    keys: Vec<Jwk>,
}

#[derive(Serialize, Clone)]
struct Jwk {
    kty: &'static str,
    kid: &'static str,
    #[serde(rename = "use")]
    use_field: &'static str,
    alg: &'static str,
    n: String,
    e: String,
}

fn base64url_no_pad(data: &[u8]) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

#[tokio::test]
async fn verify_rs256_via_jwks() {
    let private = RsaPrivateKey::new(&mut OsRng, 2048).unwrap();
    let public = private.to_public_key();

    let private_pem = private.to_pkcs1_pem(LineEnding::LF).unwrap().to_string();
    let public_pem = public.to_pkcs1_pem(LineEnding::LF).unwrap().to_string();

    let n = base64url_no_pad(&public.n().to_bytes_be());
    let e = base64url_no_pad(&public.e().to_bytes_be());

    let kid = "test-kid";
    let jwks = Jwks {
        keys: vec![Jwk {
            kty: "RSA",
            kid,
            use_field: "sig",
            alg: "RS256",
            n,
            e,
        }],
    };

    let app = Router::new().route(
        "/.well-known/jwks.json",
        get(move || {
            let jwks = jwks;
            async move { (axum::http::StatusCode::OK, Json(jwks)).into_response() }
        }),
    );

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let base_url = format!("http://{}", addr);
    let handle = tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs() as usize;

    let claims = Claims {
        sub: uuid::Uuid::new_v4().to_string(),
        tenant_id: uuid::Uuid::new_v4().to_string(),
        exp: now + 60,
        iat: now,
        iss: "iam-service".to_string(),
        roles: vec![],
        permissions: vec![],
        apps: vec![],
        apps_version: 0,
    };

    let mut header = Header::new(Algorithm::RS256);
    header.kid = Some(kid.to_string());
    let token = encode(
        &header,
        &claims,
        &EncodingKey::from_rsa_pem(private_pem.as_bytes()).unwrap(),
    )
    .unwrap();

    let cfg = JwtVerifyConfig::rs256_from_jwks(
        "iam-service",
        &format!("{}/.well-known/jwks.json", base_url),
    )
    .unwrap();
    let verified = auth_kit::jwt::verify(&token, &cfg).await.unwrap();
    assert_eq!(verified.tenant_id, claims.tenant_id);

    let cfg2 = JwtVerifyConfig::rs256_from_pem("iam-service", &public_pem).unwrap();
    let verified2 = auth_kit::jwt::verify(&token, &cfg2).await.unwrap();
    assert_eq!(verified2.sub, claims.sub);

    tokio::time::sleep(Duration::from_millis(10)).await;
    handle.abort();
}