// Minimal implementation of ECIES with x25519_dalek
use aes_gcm::{
    aead::{Aead, AeadCore},
    Aes256Gcm, Key, KeyInit, Nonce,
};
use hkdf::Hkdf;
use serde::{Deserialize, Serialize};
use sha3::Sha3_256;
use x25519_dalek::{EphemeralSecret, PublicKey, StaticSecret};

const SALT_STRING: &str = "ECIES symmetric key for Troll Patrol";

#[derive(Serialize, Deserialize)]
pub struct EciesCiphertext {
    pubkey: PublicKey,
    nonce: Vec<u8>,
    ct: Vec<u8>,
}

impl EciesCiphertext {
    pub fn encrypt(message: &[u8], receiver_pubkey: &PublicKey) -> Result<Self, String> {
        // Compute shared secret based on new ephemeral ECDH key
        let mut rng = rand::thread_rng();
        let secret = EphemeralSecret::random_from_rng(&mut rng);
        let client_pubkey = PublicKey::from(&secret);
        let shared_secret = secret.diffie_hellman(receiver_pubkey);

        // Compute key
        let hk = Hkdf::<Sha3_256>::new(None, shared_secret.as_bytes());
        let mut symmetric_key = [0u8; 32];
        if hk
            .expand(SALT_STRING.as_bytes(), &mut symmetric_key)
            .is_err()
        {
            return Err("Failed to encrypt".to_string());
        }

        // Encrypt with key
        let key: Key<Aes256Gcm> = symmetric_key.into();
        let cipher = Aes256Gcm::new(&key);
        let nonce = Aes256Gcm::generate_nonce(&mut rng);
        match cipher.encrypt(&nonce, message) {
            Ok(ct) => Ok(EciesCiphertext {
                pubkey: client_pubkey,
                nonce: nonce.to_vec(),
                ct,
            }),
            Err(_) => Err("Failed to encrypt".to_string()),
        }
    }

    pub fn decrypt(self, secret: &StaticSecret) -> Result<Vec<u8>, String> {
        // Compute shared secret
        let shared_secret = secret.diffie_hellman(&self.pubkey);

        // Compute key
        let hk = Hkdf::<Sha3_256>::new(None, shared_secret.as_bytes());
        let mut symmetric_key = [0u8; 32];
        if hk
            .expand(SALT_STRING.as_bytes(), &mut symmetric_key)
            .is_err()
        {
            return Err("Failed to decrypt".to_string());
        }

        // Decrypt with key
        let key: Key<Aes256Gcm> = symmetric_key.into();
        let cipher = Aes256Gcm::new(&key);
        let nonce = Nonce::from_slice(&self.nonce);
        match cipher.decrypt(nonce, &*self.ct) {
            Ok(m) => Ok(m),
            Err(_) => Err("Failed to decrypt".to_string()),
        }
    }
}