use std::num::NonZeroU32;

use aead::{NONCE_LEN, Nonce, UnboundKey};
use rand::SystemRandom;
use ring::{aead::{self, LessSafeKey}, digest, error::Unspecified, hmac, pbkdf2, rand::SecureRandom, rand, signature::{ KeyPair, Ed25519KeyPair, UnparsedPublicKey, ED25519 }};

const SEED_LEN: usize = 32;

fn main() -> Result<(), Unspecified> {
    {
        println!("{} HHmac {}", "-".repeat(10), "-".repeat(10));
        let rng = rand::SystemRandom::new();
        let key = hmac::Key::generate(hmac::HMAC_SHA256, &rng)?;

        let msg = "hello, world";
        let tag = hmac::sign(&key, msg.as_bytes());

        println!("{:?}", tag);
        hmac::verify(&key, msg.as_bytes(), tag.as_ref())?;
        println!("Verify success");
    }
    {
        println!("{} SHA256 {}", "-".repeat(10), "-".repeat(10));
        let sha256 = digest::digest(&digest::SHA256, b"hello, world");
        println!("{:?}", sha256);
    }
    {
        println!("{} EdDSA  {}", "-".repeat(10), "-".repeat(10));
        let rng = rand::SystemRandom::new();
        let seed: [u8; SEED_LEN] = rand::generate(&rng)?.expose();
        println!("seed: {}", hex::encode(&seed));
        let key_pair = Ed25519KeyPair::from_seed_unchecked(&seed)?;
        println!("private: {:?}", key_pair);
        let sig = key_pair.sign(&"hello world".as_bytes());
        println!("signature: {}", hex::encode(sig.as_ref()));
        let public_key = key_pair.public_key().as_ref();
        let verify_result = UnparsedPublicKey::new(&ED25519, &public_key).verify(&"hello world".as_bytes(), sig.as_ref());
        println!("verify: {}", verify_result.is_ok());
    }
    {
        println!("{} Random  {}", "-".repeat(10), "-".repeat(10));
        let mut random_bytes = [0_u8; 32];
        let random = SystemRandom::new();
        random.fill(&mut random_bytes).unwrap();
        println!("Random bytes: {}", hex::encode(&random_bytes));
    }
    {
        println!("{} Aead Crypto  {}", "-".repeat(10), "-".repeat(10));
        let mut key = vec![];
        for _ in 0..16 { key.push(0_u8); }
        let aead_crypto = AeadCrypto::new(&key, &aead::AES_128_GCM);
        let mut text = "hello world".as_bytes().to_vec();
        println!("Cleartext: {}", hex::encode(&text));
        aead_crypto.encrypt(&mut text);
        println!("Encrypted: {}", hex::encode(&text));
        let decrypt_result = aead_crypto.decrypt(&mut text);
        println!("Decrypted: {}, text: {}", decrypt_result, hex::encode(text));
    }

    Ok(())
}

pub trait Crypto: Send + Sync {
    fn encrypt(&self, buf: &mut Vec<u8>);
    fn decrypt(&self, buf: &mut Vec<u8>) -> bool;
}

pub struct AeadCrypto {
    key: LessSafeKey,
    algorithm: &'static aead::Algorithm,
    random: SystemRandom,
}

impl AeadCrypto {
    pub fn new(key: &[u8], algorithm: &'static aead::Algorithm) -> Self {
        let salt = b"ap-kcp-aead-salt";
        let mut key_bytes = Vec::with_capacity(algorithm.key_len());
        key_bytes.resize(algorithm.key_len(), 0);
        pbkdf2::derive(
            pbkdf2::PBKDF2_HMAC_SHA256,
            NonZeroU32::new(32).unwrap(),
            salt,
            key,
            &mut key_bytes,
        );
        let key = UnboundKey::new(algorithm, &key_bytes).unwrap();
        let key = LessSafeKey::new(key);
        Self {
            key,
            algorithm,
            random: SystemRandom::new(),
        }
    }
}

impl Crypto for AeadCrypto {
    fn encrypt(&self, buf: &mut Vec<u8>) {
        let mut nonce_bytes = [0u8; NONCE_LEN];
        self.random.fill(&mut nonce_bytes).unwrap();
        let nonce = Nonce::assume_unique_for_key(nonce_bytes);

        // | ENCRPYTED | TAG | NONCE |
        self.key
            .seal_in_place_append_tag(nonce, aead::Aad::empty(), buf)
            .unwrap();

        buf.extend_from_slice(&nonce_bytes);
    }

    fn decrypt(&self, buf: &mut Vec<u8>) -> bool {
        if buf.len() < aead::NONCE_LEN + self.algorithm.tag_len() {
            return false;
        }

        let len = buf.len();
        let mut nonce_bytes = [0u8; aead::NONCE_LEN];
        nonce_bytes.copy_from_slice(&buf[len - aead::NONCE_LEN..]);
        let nonce = Nonce::assume_unique_for_key(nonce_bytes);

        buf.truncate(len - aead::NONCE_LEN);

        let len = if let Ok(plaintext) = self.key.open_in_place(nonce, aead::Aad::empty(), buf) {
            plaintext.len()
        } else {
            println!("[ERROR] failed to decrypt");
            return false;
        };
        buf.truncate(len);
        true
    }
}