use josekit::jwe;
use josekit::jwe::alg::aeskw::AeskwJweAlgorithm;
use josekit::jwe::alg::rsaes::RsaesJweAlgorithm;
use josekit::jwe::JweHeader;
use josekit::jwk::alg::rsa::RsaKeyPair;
use josekit::jwk::Jwk;
use rust_util::XResult;
use serde_json::Value;

const LOCAL_KMS_PREFIX: &'static str = "LKMS:";

pub fn generate_rsa_key(bits: u32) -> XResult<RsaKeyPair> {
    Ok(RsaKeyPair::generate(bits)?)
}

// pub fn serialize_jwe_rsa(payload: &[u8], jwk: &Jwk) -> XResult<String> {
//     let mut header = JweHeader::new();
//     header.set_content_encryption("A256GCM");
//     header.set_claim("vendor", Some(Value::String("local-mini-kms".to_string())))?;
//     let encrypter = RsaesJweAlgorithm::RsaOaep.encrypter_from_jwk(&jwk)?;
//     Ok(format!("{}{}", LOCAL_KMS_PREFIX, jwe::serialize_compact(payload, &header, &encrypter)?))
// }

pub fn deserialize_jwe_rsa(jwe: &str, jwk: &Jwk) -> XResult<(Vec<u8>, JweHeader)> {
    let decrypter = RsaesJweAlgorithm::RsaOaep.decrypter_from_jwk(jwk)?;
    Ok(jwe::deserialize_json(&get_jwe(jwe), &decrypter)?)
}

pub fn serialize_jwe_aes(payload: &[u8], key: &[u8]) -> XResult<String> {
    let mut header = JweHeader::new();
    header.set_content_encryption("A256GCM");
    header.set_claim("vendor", Some(Value::String("local-mini-kms".to_string())))?;
    let encrypter = AeskwJweAlgorithm::A256kw.encrypter_from_bytes(key)?;
    Ok(format!("{}{}", LOCAL_KMS_PREFIX, jwe::serialize_compact(payload, &header, &encrypter)?))
}

pub fn deserialize_jwe_aes(jwe: &str, key: &[u8]) -> XResult<(Vec<u8>, JweHeader)> {
    let decrypter = AeskwJweAlgorithm::A256kw.decrypter_from_bytes(key)?;
    Ok(jwe::deserialize_compact(&get_jwe(jwe), &decrypter)?)
}

fn get_jwe(jwe: &str) -> String {
    if jwe.starts_with(LOCAL_KMS_PREFIX) {
        jwe.chars().skip(LOCAL_KMS_PREFIX.len()).collect()
    } else {
        jwe.to_string()
    }
}