use aes_gcm::aead::{Aead, KeyInit};
use aes_gcm::{Aes256Gcm, Nonce};
use falcon::FnDsaKeyPair;
use pbkdf2::pbkdf2_hmac;
use rand::{rngs::OsRng, RngCore};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use zeroize::Zeroize;

use crate::address;
use crate::crypto;
use crate::error::BrowserCoreError;
use crate::legacy_wallet;
use crate::{PRIVATE_KEY_LENGTH, PUBLIC_KEY_LENGTH};

const KDF_ITERATIONS: u32 = 310_000;
const REGISTRATION_COMMAND: u8 = 38;

#[derive(Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct BrowserWallet {
    pub version: u8,
    pub network: String,
    pub short_address: String,
    pub vanity_address: Option<String>,
    pub public_key: String,
    pub salt: String,
    pub iterations: u32,
    pub nonce: String,
    pub encrypted_private_key: String,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct LegacySavedWallet {
    pub short_address: String,
    pub vanity_address: Option<String>,
    pub public_key: String,
    pub private_key: String,
}

#[derive(Debug, Serialize)]
pub struct UnlockedWallet {
    pub short_address: String,
    pub vanity_address: Option<String>,
    pub public_key: String,
    pub private_key: String,
}

#[derive(Debug, Serialize)]
pub struct RegistrationPayload {
    pub command: u8,
    pub short_address: String,
    pub public_key: String,
    pub hash: String,
    pub signature: String,
    pub payload_hex: String,
}

pub fn generate(password: &str) -> Result<BrowserWallet, BrowserCoreError> {
    let (public_key, mut private_key) = generate_keypair()?;
    let wallet = encrypt_wallet(&public_key, &private_key, None, password);
    private_key.zeroize();
    wallet
}

pub fn unlock(wallet: &BrowserWallet, password: &str) -> Result<UnlockedWallet, BrowserCoreError> {
    if wallet.version != 1 || wallet.iterations < 100_000 {
        return Err(BrowserCoreError::InvalidRequest(
            "unsupported browser wallet format".into(),
        ));
    }
    let salt = decode_length(&wallet.salt, 32, "wallet salt")?;
    let nonce = decode_length(&wallet.nonce, 12, "wallet nonce")?;
    let ciphertext = hex::decode(&wallet.encrypted_private_key)
        .map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    let key = derive_key(password, &salt, wallet.iterations);
    let private_key = Aes256Gcm::new_from_slice(&key)
        .unwrap()
        .decrypt(Nonce::from_slice(&nonce), ciphertext.as_ref())
        .map_err(|_| {
            BrowserCoreError::InvalidRequest("wallet password or integrity check failed".into())
        })?;
    validate_key_material(&private_key, &wallet.public_key, &wallet.short_address)?;
    Ok(UnlockedWallet {
        short_address: wallet.short_address.clone(),
        vanity_address: wallet.vanity_address.clone(),
        public_key: wallet.public_key.clone(),
        private_key: hex::encode(private_key),
    })
}

pub fn import_legacy(
    legacy: &LegacySavedWallet,
    legacy_password: &str,
    browser_password: &str,
) -> Result<BrowserWallet, BrowserCoreError> {
    let private_key_hex = legacy_wallet::decrypt_private_key(&legacy.private_key, legacy_password)?;
    let mut private_key =
        hex::decode(&private_key_hex).map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    validate_key_material(&private_key, &legacy.public_key, &legacy.short_address)?;
    let public_key =
        hex::decode(&legacy.public_key).map_err(|_| BrowserCoreError::InvalidPublicKey)?;
    let wallet = encrypt_wallet(
        &public_key,
        &private_key,
        legacy.vanity_address.clone(),
        browser_password,
    );
    private_key.zeroize();
    wallet
}

pub fn import_private_key_image(
    image_base64: &str,
    image_password: &str,
    browser_password: &str,
) -> Result<BrowserWallet, BrowserCoreError> {
    let mut private_key_hex = legacy_wallet::decrypt_private_key(image_base64, image_password)?;
    let result = import_private_key(&private_key_hex, browser_password);
    private_key_hex.zeroize();
    result
}

pub fn import_private_key(
    private_key_hex: &str,
    browser_password: &str,
) -> Result<BrowserWallet, BrowserCoreError> {
    let mut private_key =
        hex::decode(private_key_hex).map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    if private_key.len() != PRIVATE_KEY_LENGTH {
        private_key.zeroize();
        return Err(BrowserCoreError::InvalidPrivateKey);
    }
    let result = wallet_from_private_key(&private_key, browser_password);
    private_key.zeroize();
    result
}

pub fn export_legacy(
    wallet: &BrowserWallet,
    browser_password: &str,
    legacy_password: &str,
) -> Result<LegacySavedWallet, BrowserCoreError> {
    let mut unlocked = unlock(wallet, browser_password)?;
    let image = legacy_wallet::encrypt_private_key(&unlocked.private_key, legacy_password)?;
    unlocked.private_key.zeroize();
    Ok(LegacySavedWallet {
        short_address: wallet.short_address.clone(),
        vanity_address: wallet.vanity_address.clone(),
        public_key: wallet.public_key.clone(),
        private_key: image,
    })
}

pub fn registration(
    wallet: &BrowserWallet,
    password: &str,
) -> Result<RegistrationPayload, BrowserCoreError> {
    let mut unlocked = unlock(wallet, password)?;
    let result = registration_from_keys(
        &unlocked.short_address,
        &unlocked.public_key,
        &unlocked.private_key,
    );
    unlocked.private_key.zeroize();
    result
}

pub fn registration_from_keys(
    short_address: &str,
    public_key_hex: &str,
    private_key_hex: &str,
) -> Result<RegistrationPayload, BrowserCoreError> {
    let address_bytes = address::short_address_to_bytes(short_address)?;
    let public_key = decode_length(public_key_hex, PUBLIC_KEY_LENGTH, "public key")?;
    let mut signed_payload = Vec::with_capacity(1 + 22 + PUBLIC_KEY_LENGTH);
    signed_payload.push(REGISTRATION_COMMAND);
    signed_payload.extend_from_slice(&address_bytes);
    signed_payload.extend_from_slice(&public_key);
    let hash = crypto::skein_256(&signed_payload);
    let private_key =
        hex::decode(private_key_hex).map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    validate_key_material(&private_key, public_key_hex, short_address)?;
    let signature = crypto::sign(&hash, &private_key)?;
    let mut payload = signed_payload.clone();
    payload.extend_from_slice(&signature);
    Ok(RegistrationPayload {
        command: REGISTRATION_COMMAND,
        short_address: short_address.to_owned(),
        public_key: public_key_hex.to_owned(),
        hash: hex::encode(hash),
        signature: hex::encode(signature),
        payload_hex: hex::encode(payload),
    })
}

pub fn handshake_proof(
    public_key_hex: &str,
    private_key_hex: &str,
) -> Result<String, BrowserCoreError> {
    let public_key = decode_length(public_key_hex, PUBLIC_KEY_LENGTH, "public key")?;
    let private_key =
        hex::decode(private_key_hex).map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    let keypair = FnDsaKeyPair::from_private_key(&private_key)
        .map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    if keypair.public_key() != public_key {
        return Err(BrowserCoreError::InvalidRequest(
            "wallet public key does not match its private key".into(),
        ));
    }
    Ok(hex::encode(crypto::sign(
        &crypto::skein_256(b"aced"),
        &private_key,
    )?))
}

fn generate_keypair() -> Result<(Vec<u8>, Vec<u8>), BrowserCoreError> {
    loop {
        let keypair = FnDsaKeyPair::generate(9).map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
        let public_key = keypair.public_key().to_vec();
        if address::valid_public_key(&public_key) {
            return Ok((public_key, keypair.private_key().to_vec()));
        }
    }
}

fn wallet_from_private_key(
    private_key: &[u8],
    password: &str,
) -> Result<BrowserWallet, BrowserCoreError> {
    let keypair = FnDsaKeyPair::from_private_key(private_key)
        .map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    let public_key = keypair.public_key();
    address::public_key_to_short_address(public_key)?;
    encrypt_wallet(public_key, private_key, None, password)
}

fn encrypt_wallet(
    public_key: &[u8],
    private_key: &[u8],
    vanity_address: Option<String>,
    password: &str,
) -> Result<BrowserWallet, BrowserCoreError> {
    let short_address = address::public_key_to_short_address(public_key)?;
    let mut salt = [0u8; 32];
    let mut nonce = [0u8; 12];
    OsRng.fill_bytes(&mut salt);
    OsRng.fill_bytes(&mut nonce);
    let mut key = derive_key(password, &salt, KDF_ITERATIONS);
    let ciphertext = Aes256Gcm::new_from_slice(&key)
        .unwrap()
        .encrypt(Nonce::from_slice(&nonce), private_key)
        .map_err(|_| BrowserCoreError::InvalidRequest("wallet encryption failed".into()))?;
    key.zeroize();
    Ok(BrowserWallet {
        version: 1,
        network: address::network_suffix().into(),
        short_address,
        vanity_address,
        public_key: hex::encode(public_key),
        salt: hex::encode(salt),
        iterations: KDF_ITERATIONS,
        nonce: hex::encode(nonce),
        encrypted_private_key: hex::encode(ciphertext),
    })
}

fn derive_key(password: &str, salt: &[u8], iterations: u32) -> [u8; 32] {
    let mut key = [0u8; 32];
    pbkdf2_hmac::<Sha256>(password.as_bytes(), salt, iterations, &mut key);
    key
}

fn validate_key_material(
    private_key: &[u8],
    public_key_hex: &str,
    short_address: &str,
) -> Result<(), BrowserCoreError> {
    if private_key.len() != PRIVATE_KEY_LENGTH {
        return Err(BrowserCoreError::InvalidPrivateKey);
    }
    let keypair = FnDsaKeyPair::from_private_key(private_key)
        .map_err(|_| BrowserCoreError::InvalidPrivateKey)?;
    let public_key = keypair.public_key();
    if public_key.len() != PUBLIC_KEY_LENGTH || hex::encode(public_key) != public_key_hex {
        return Err(BrowserCoreError::InvalidRequest(
            "wallet public key does not match its private key".into(),
        ));
    }
    if address::public_key_to_short_address(public_key)? != short_address {
        return Err(BrowserCoreError::InvalidRequest(
            "wallet address does not match its private key".into(),
        ));
    }
    Ok(())
}

fn decode_length(value: &str, length: usize, name: &str) -> Result<Vec<u8>, BrowserCoreError> {
    let bytes = hex::decode(value)
        .map_err(|_| BrowserCoreError::InvalidRequest(format!("{name} is not hexadecimal")))?;
    if bytes.len() != length {
        return Err(BrowserCoreError::InvalidRequest(format!(
            "{name} has an invalid length"
        )));
    }
    Ok(bytes)
}

#[cfg(test)]
mod tests {
    use super::{export_legacy, generate, import_private_key, import_private_key_image, unlock};

    #[test]
    fn private_key_import_recreates_the_same_wallet() {
        let original = generate("original password").unwrap();
        let unlocked = unlock(&original, "original password").unwrap();
        let imported = import_private_key(&unlocked.private_key, "new password").unwrap();
        let imported_unlocked = unlock(&imported, "new password").unwrap();

        assert_eq!(imported.short_address, original.short_address);
        assert_eq!(imported.public_key, original.public_key);
        assert_eq!(imported_unlocked.private_key, unlocked.private_key);
    }

    #[test]
    fn private_key_image_import_recreates_the_same_wallet() {
        let original = generate("browser password").unwrap();
        let legacy = export_legacy(&original, "browser password", "image password").unwrap();
        let imported =
            import_private_key_image(&legacy.private_key, "image password", "new password")
                .unwrap();

        assert_eq!(imported.short_address, original.short_address);
        assert_eq!(imported.public_key, original.public_key);
        assert!(unlock(&imported, "new password").is_ok());
    }
}
