use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Validation};
use ring::signature::{KeyPair, UnparsedPublicKey};
use serde::de::DeserializeOwned;
use serde::Serialize;

#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
#[serde(tag = "alg")]
pub enum TokenPubKey {
    /// This variant DOES make crash the program. It MUST NOT be serialized.
    ///
    /// It is a hack to hide public key when getting the list of tokens
    None,

    /// ECDSA with SHA2-384 variant
    ES384 { r#pub: String },
}

#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
#[serde(tag = "alg")]
pub enum TokenPrivKey {
    ES384 { r#priv: String },
}

/// Generate a new token keypair
pub fn generate_key_pair() -> anyhow::Result<(TokenPubKey, TokenPrivKey)> {
    let doc = ring::signature::EcdsaKeyPair::generate_pkcs8(
        &ring::signature::ECDSA_P384_SHA384_ASN1_SIGNING,
        &ring::rand::SystemRandom::new(),
    )?;

    let priv_pem = pem::encode(&pem::Pem::new("PRIVATE KEY", doc.as_ref()));

    let pair = ring::signature::EcdsaKeyPair::from_pkcs8(
        &ring::signature::ECDSA_P384_SHA384_ASN1_SIGNING,
        doc.as_ref(),
        &ring::rand::SystemRandom::new(),
    )?;
    let pub_pem = pem::encode(&pem::Pem::new("PUBLIC KEY", pair.public_key().as_ref()));


    let pk = pair.public_key();
    let unp = UnparsedPublicKey::new(&ring::signature::ECDSA_P384_SHA384_ASN1_SIGNING, pk.as_ref());

    let decoding_key = DecodingKey::from_ec_pem(pub_pem.as_bytes()).expect("aie ai");

    Ok((
        TokenPubKey::ES384 { r#pub: pub_pem },
        TokenPrivKey::ES384 { r#priv: priv_pem },
    ))
}

/// Sign JWT with a private key
pub fn sign_jwt<C: Serialize>(key: &TokenPrivKey, claims: &C) -> anyhow::Result<String> {
    match key {
        TokenPrivKey::ES384 { r#priv } => {
            let encoding_key = EncodingKey::from_ec_pem(r#priv.as_bytes())?;

            Ok(jsonwebtoken::encode(
                &jsonwebtoken::Header::new(Algorithm::ES384),
                &claims,
                &encoding_key,
            )?)
        }
    }
}

/// Validate a given JWT
pub fn validate_jwt<E: DeserializeOwned>(key: &TokenPubKey, token: &str) -> anyhow::Result<E> {
    match key {
        TokenPubKey::ES384 { r#pub } => {
            let decoding_key = DecodingKey::from_ec_pem(r#pub.as_bytes())?;

            let validation = Validation::new(Algorithm::ES384);
            Ok(jsonwebtoken::decode::<E>(token, &decoding_key, &validation)?.claims)
        }
        TokenPubKey::None => {
            panic!("A public key is required!")
        }
    }
}

#[cfg(test)]
mod test {
    use crate::utils::jwt_utils::{generate_key_pair, sign_jwt, validate_jwt};
    use crate::utils::time_utils::time;
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
    pub struct Claims {
        sub: String,
        exp: u64,
    }

    #[test]
    fn jwt_encode_sign_verify_valid() {
        let (pub_key, priv_key) = generate_key_pair().unwrap();
        let claims = Claims {
            sub: "my-sub".to_string(),
            exp: time() + 100,
        };
        let jwt = sign_jwt(&priv_key, &claims).expect("Failed to sign JWT!");

        println!("pub {pub_key:?}");
        println!("priv {priv_key:?}");
        let claims_out = validate_jwt(&pub_key, &jwt).expect("Failed to validate JWT!");

        assert_eq!(claims, claims_out)
    }
}