use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use rand_chacha::ChaCha20Rng;
use rand_core::SeedableRng;
use rsa::pkcs1::{EncodeRsaPrivateKey, EncodeRsaPublicKey};
use rsa::traits::PublicKeyParts;
use rsa::RsaPrivateKey;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;

#[derive(Clone)]
pub struct CryptoProvider {
    encoding_key: EncodingKey,
    decoding_key: DecodingKey,
    jwks_json: serde_json::Value,
    kid: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct AccessTokenClaims {
    pub sub: String,
    pub iss: String,
    pub client_id: String,
    pub token_use: String,
    pub scope: String,
    pub auth_time: u64,
    pub exp: u64,
    pub iat: u64,
    pub jti: String,
    pub username: String,
    #[serde(rename = "cognito:groups", skip_serializing_if = "Vec::is_empty")]
    pub cognito_groups: Vec<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct IdTokenClaims {
    pub sub: String,
    pub iss: String,
    pub aud: String,
    pub token_use: String,
    pub auth_time: u64,
    pub exp: u64,
    pub iat: u64,
    pub jti: String,
    #[serde(rename = "cognito:username")]
    pub cognito_username: String,
    #[serde(rename = "cognito:groups", skip_serializing_if = "Vec::is_empty")]
    pub cognito_groups: Vec<String>,
    #[serde(flatten)]
    pub attributes: HashMap<String, serde_json::Value>,
}

impl CryptoProvider {
    pub fn new(signing_secret: &str) -> Self {
        let seed_hash = Sha256::digest(signing_secret.as_bytes());
        let mut seed = [0u8; 32];
        seed.copy_from_slice(&seed_hash);
        let mut rng = ChaCha20Rng::from_seed(seed);

        let private_key =
            RsaPrivateKey::new(&mut rng, 2048).expect("failed to generate RSA key");

        let private_key_der = private_key
            .to_pkcs1_der()
            .expect("failed to encode private key");
        let public_key_der = private_key
            .to_public_key()
            .to_pkcs1_der()
            .expect("failed to encode public key");

        let encoding_key = EncodingKey::from_rsa_der(private_key_der.as_bytes());
        let decoding_key = DecodingKey::from_rsa_der(public_key_der.as_bytes());

        let kid_hash = Sha256::digest(public_key_der.as_bytes());
        let kid = URL_SAFE_NO_PAD.encode(&kid_hash[..8]);

        let n = private_key.n();
        let e = private_key.e();
        let n_bytes = n.to_bytes_be();
        let e_bytes = e.to_bytes_be();

        let jwks_json = serde_json::json!({
            "keys": [{
                "kty": "RSA",
                "alg": "RS256",
                "use": "sig",
                "kid": kid,
                "n": URL_SAFE_NO_PAD.encode(&n_bytes),
                "e": URL_SAFE_NO_PAD.encode(&e_bytes),
            }]
        });

        Self {
            encoding_key,
            decoding_key,
            jwks_json,
            kid,
        }
    }

    pub fn jwks(&self) -> &serde_json::Value {
        &self.jwks_json
    }

    fn now_secs() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }

    pub fn sign_access_token(
        &self,
        sub: &str,
        username: &str,
        _pool_id: &str,
        client_id: &str,
        issuer: &str,
        groups: Vec<String>,
    ) -> Result<String, jsonwebtoken::errors::Error> {
        let now = Self::now_secs();
        let claims = AccessTokenClaims {
            sub: sub.to_string(),
            iss: issuer.to_string(),
            client_id: client_id.to_string(),
            token_use: "access".to_string(),
            scope: "aws.cognito.signin.user.admin".to_string(),
            auth_time: now,
            exp: now + 3600,
            iat: now,
            jti: Uuid::new_v4().to_string(),
            username: username.to_string(),
            cognito_groups: groups,
        };
        let mut header = Header::new(Algorithm::RS256);
        header.kid = Some(self.kid.clone());
        encode(&header, &claims, &self.encoding_key)
    }

    pub fn sign_id_token(
        &self,
        sub: &str,
        username: &str,
        _pool_id: &str,
        client_id: &str,
        issuer: &str,
        groups: Vec<String>,
        attributes: HashMap<String, serde_json::Value>,
    ) -> Result<String, jsonwebtoken::errors::Error> {
        let now = Self::now_secs();
        let claims = IdTokenClaims {
            sub: sub.to_string(),
            iss: issuer.to_string(),
            aud: client_id.to_string(),
            token_use: "id".to_string(),
            auth_time: now,
            exp: now + 3600,
            iat: now,
            jti: Uuid::new_v4().to_string(),
            cognito_username: username.to_string(),
            cognito_groups: groups,
            attributes,
        };
        let mut header = Header::new(Algorithm::RS256);
        header.kid = Some(self.kid.clone());
        encode(&header, &claims, &self.encoding_key)
    }

    pub fn decode_access_token(
        &self,
        token: &str,
    ) -> Result<AccessTokenClaims, jsonwebtoken::errors::Error> {
        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_required_spec_claims(&["sub", "exp", "iss"]);
        validation.validate_exp = true;
        validation.validate_aud = false;
        let data = decode::<AccessTokenClaims>(token, &self.decoding_key, &validation)?;
        Ok(data.claims)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_deterministic_key_generation() {
        let cp1 = CryptoProvider::new("test-secret");
        let cp2 = CryptoProvider::new("test-secret");
        assert_eq!(cp1.kid, cp2.kid);
        assert_eq!(cp1.jwks_json, cp2.jwks_json);
    }

    #[test]
    fn test_different_secrets_produce_different_keys() {
        let cp1 = CryptoProvider::new("secret-a");
        let cp2 = CryptoProvider::new("secret-b");
        assert_ne!(cp1.kid, cp2.kid);
    }

    #[test]
    fn test_sign_and_decode_access_token() {
        let cp = CryptoProvider::new("test-secret");
        let token = cp
            .sign_access_token(
                "sub-123",
                "testuser",
                "us-east-1_abc",
                "client-id",
                "http://localhost:9229/us-east-1_abc",
                vec!["admin".to_string()],
            )
            .unwrap();

        let claims = cp.decode_access_token(&token).unwrap();
        assert_eq!(claims.sub, "sub-123");
        assert_eq!(claims.username, "testuser");
        assert_eq!(claims.token_use, "access");
        assert_eq!(claims.cognito_groups, vec!["admin"]);
    }

    #[test]
    fn test_jwks_structure() {
        let cp = CryptoProvider::new("test-secret");
        let jwks = cp.jwks();
        let keys = jwks["keys"].as_array().unwrap();
        assert_eq!(keys.len(), 1);
        assert_eq!(keys[0]["kty"], "RSA");
        assert_eq!(keys[0]["alg"], "RS256");
        assert_eq!(keys[0]["use"], "sig");
        assert!(keys[0]["kid"].is_string());
        assert!(keys[0]["n"].is_string());
        assert!(keys[0]["e"].is_string());
    }
}
