use std::collections::HashMap;
use std::sync::LazyLock;

use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, encode};
use rcgen::{KeyPair, PKCS_RSA_SHA256, RsaKeySize};
use remora_domain::{BucketName, StorageScope, TargetPrefix, TargetStoreId};
use remora_identity::{
    OidcConfig, OidcIdentityProvider, ZITADEL_PROJECT_ROLES_CLAIM, encode_admin_role,
    encode_storage_role,
};
use serde_json::json;

struct EphemeralRsa {
    encoding: EncodingKey,
    decoding: DecodingKey,
}

static EPHEMERAL_RSA: LazyLock<EphemeralRsa> = LazyLock::new(|| {
    let key_pair = KeyPair::generate_rsa_for(&PKCS_RSA_SHA256, RsaKeySize::_2048)
        .expect("generate ephemeral RSA key");
    let private_pem = key_pair.serialize_pem();
    let public_pem = key_pair.public_key_pem();
    EphemeralRsa {
        encoding: EncodingKey::from_rsa_pem(private_pem.as_bytes()).expect("encoding key"),
        decoding: DecodingKey::from_rsa_pem(public_pem.as_bytes()).expect("decoding key"),
    }
});

fn storage_scope() -> StorageScope {
    StorageScope {
        store: TargetStoreId::parse("target-main").unwrap(),
        bucket: BucketName::parse("target-bucket").unwrap(),
        prefix: TargetPrefix::parse("models", []).unwrap(),
    }
}

fn provider(with_zitadel_claim: bool) -> OidcIdentityProvider {
    let mut keys = HashMap::new();
    keys.insert("test-kid".into(), EPHEMERAL_RSA.decoding.clone());
    let mut config = OidcConfig::new("https://issuer.example", "remora-client").expect("config");
    if with_zitadel_claim {
        config = config.with_zitadel_project_roles();
    }
    OidcIdentityProvider::with_static_keys(config, keys).expect("provider")
}

fn mint(claims: serde_json::Value) -> String {
    let mut header = Header::new(Algorithm::RS256);
    header.kid = Some("test-kid".into());
    encode(&header, &claims, &EPHEMERAL_RSA.encoding).expect("encode")
}

#[tokio::test]
async fn valid_token_maps_roles_from_configured_zitadel_claim() {
    let storage = encode_storage_role(&storage_scope()).unwrap();
    let admin = encode_admin_role();
    let mut project_roles = serde_json::Map::new();
    project_roles.insert(admin, json!({ "org-1": "example.local" }));
    project_roles.insert(storage, json!({ "org-1": "example.local" }));
    let token = mint(json!({
        "iss": "https://issuer.example",
        "sub": "user-12345678",
        "aud": "remora-client",
        "exp": 4_102_444_800_i64,
        ZITADEL_PROJECT_ROLES_CLAIM: project_roles,
        // Untrusted sources must not grant authority.
        "realm_access": { "roles": ["admin", "remora:v1:admin"] },
        "roles": ["remora:v1:admin"],
    }));
    let identity = provider(true)
        .authenticate_token(&token)
        .await
        .expect("valid token");
    assert_eq!(identity.subject, "user-12345678");
    assert!(identity.roles.administrator);
    assert_eq!(identity.roles.scopes.len(), 1);
    assert_eq!(identity.roles.scopes[0].prefix.canonical(), "models");
}

#[tokio::test]
async fn unconfigured_project_claim_and_bare_admin_do_not_grant() {
    let mut project_roles = serde_json::Map::new();
    project_roles.insert(encode_admin_role(), json!({ "org-1": "example.local" }));
    let token = mint(json!({
        "iss": "https://issuer.example",
        "sub": "user-12345678",
        "aud": "remora-client",
        "exp": 4_102_444_800_i64,
        "urn:zitadel:iam:org:project:other:roles": project_roles,
        "resource_access": {
            "evil-client": { "roles": ["remora:v1:admin"] }
        },
        "realm_access": { "roles": ["admin"] },
    }));
    let identity = provider(true)
        .authenticate_token(&token)
        .await
        .expect("valid token without remora grants");
    assert!(!identity.roles.administrator);
    assert!(identity.roles.scopes.is_empty());
}

#[tokio::test]
async fn keycloak_resource_access_client_roles_decode() {
    let token = mint(json!({
        "iss": "https://issuer.example",
        "sub": "kc-user-0001",
        "aud": ["remora-client", "account"],
        "exp": 4_102_444_800_i64,
        "resource_access": {
            "remora-client": { "roles": [encode_admin_role()] },
            "account": { "roles": ["manage-account"] }
        },
    }));
    let identity = provider(false)
        .authenticate_token(&token)
        .await
        .expect("valid keycloak-shaped token");
    assert!(identity.roles.administrator);
}

#[tokio::test]
async fn invalid_issuer_audience_exp_and_signature_fail_closed() {
    let good_base = json!({
        "iss": "https://issuer.example",
        "sub": "user-12345678",
        "aud": "remora-client",
        "exp": 4_102_444_800_i64,
        "resource_access": {
            "remora-client": { "roles": [encode_admin_role()] }
        },
    });

    let wrong_issuer = {
        let mut claims = good_base.clone();
        claims["iss"] = json!("https://evil.example");
        mint(claims)
    };
    let wrong_aud = {
        let mut claims = good_base.clone();
        claims["aud"] = json!("other-client");
        mint(claims)
    };
    let expired = {
        let mut claims = good_base.clone();
        claims["exp"] = json!(1_000_000_i64);
        mint(claims)
    };

    for token in [wrong_issuer, wrong_aud, expired, "totally-not-a-jwt".into()] {
        let result = provider(false).authenticate_token(&token).await;
        assert!(
            matches!(result, Err(remora_identity::IdentityError::InvalidToken)),
            "token should be rejected: {result:?}"
        );
    }
}