use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use sha2::Digest;
use smcrypto::{sm2, sm3, sm4};
use crate::error::Sm2Error;

/// # 生成密钥对
///
/// - 返回值：(私钥，公钥)
pub fn generate_keypair() -> (String, String) {
    let (sk, pk) = sm2::gen_keypair();
    (sk, format!("04{}", pk))
}

/// # 加密
///
/// - data: 待加密数据
/// - secret: 私钥（ APP唯一秘钥：X_SPDB_CLIENT_SECRET, 秘钥有过期时间，注意存储方式 ）
pub fn encrypt(data: &str, secret: &str) -> Result<String, Sm2Error> {
    let key = hex::decode(md5_key(secret))?;
    let sm4_ecb = sm4::CryptSM4ECB::new(&key);
    let encrypted = sm4_ecb.encrypt_ecb(data.as_bytes());
    Ok(STANDARD.encode(hex::encode(encrypted)))
}

/// # 解密
///
/// - encrypt_data: 加密数据
/// - secret: 私钥（ APP唯一秘钥：X_SPDB_CLIENT_SECRET, 秘钥有过期时间，注意存储方式 ）
pub fn decrypt(encrypted: &str, secret: &str) -> Result<String, Sm2Error> {
    let key = hex::decode(md5_key(secret))?;
    let encrypt_data = hex::decode(STANDARD.decode(encrypted)?)?;
    let sm4_ecb = sm4::CryptSM4ECB::new(&key);
    let decrypted = sm4_ecb.decrypt_ecb(&encrypt_data);
    Ok(String::from_utf8(decrypted)?)
}

/// # md5生成key
///
/// - 通过SHA-256生成X-SPDB-Client-Secret转16进制的摘要，得到sha256Secret
/// - sm3做hash处理sha256Secret
/// - 通过md5生成sm3Hex转16进制的32位摘要作为加解密key
pub fn md5_key(secret: &str) -> String {
    let mut sha2 = sha2::Sha256::new();
    sha2.update(secret);
    let hash = hex::encode(sha2.finalize());
    let hash = sm3::sm3_hash(hash.as_bytes());
    format!("{:x}", md5::compute(hash))
}

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

    #[test]
    fn test_generate_keypair() {
        let (sk, pk) = generate_keypair();
        println!("sk: {}, pk: {}", sk, pk);
    }

    #[test]
    fn test_encrypt() {
        let msg = r#"hello world"#;
        let secret = "ZTPkMS00ZNTP5NzPwNjAu";

        // 加密
        let encrypt_data = encrypt(msg, secret).unwrap();
        assert_eq!(encrypt_data, "M2YyZWFlOTU4MzBkZTUxMGQyOTNjNmUzYzA1ODg2NjM=");

        // 解密
        let decrypt_data = decrypt(&encrypt_data, secret).unwrap();
        assert_eq!(decrypt_data, msg);
    }
}