use crate::binary;

use aes::{block_cipher_trait::generic_array::GenericArray, Aes256};
use block_modes::{block_padding::ZeroPadding, BlockMode, Ecb};
use hmac::crypto_mac::MacResult;
use hmac::{Hmac, Mac};
use sha2::{Digest, Sha256, Sha512};
use std::string::ToString;
use thiserror::Error;
type HmacSha256 = Hmac<Sha256>;

/// Credentials needed to unlock the database
///
/// Currently it supports unlocking a database with a combination
/// of password, keyfile or both.
///
/// For the compmon case of creating credentials from just a password,
/// you can use
///
/// ```
/// # use kdbx_rs::CompositeKey;
/// CompositeKey::from_password("abcdef");
/// ```
///
/// Otherwise you can use [`CompositeKey::new`] to provide other combinations
pub struct CompositeKey {
    pw: Option<String>,
    keyfile: Option<Vec<u8>>,
}

impl CompositeKey {
    /// Create a new composite key
    pub fn new(pw: Option<String>, keyfile: Option<Vec<u8>>) -> CompositeKey {
        CompositeKey { pw, keyfile }
    }

    /// Utility method for making a key with just a password
    pub fn from_password(pw: &str) -> CompositeKey {
        CompositeKey::new(Some(pw.into()), None)
    }

    pub(crate) fn composed(&self) -> ComposedKey {
        let mut buffer = Vec::new();
        if let Some(ref pw) = self.pw {
            buffer.extend(Sha256::digest(pw.as_bytes()))
            //buffer.extend(pw.as_bytes());
        }
        if let Some(ref keyfile) = self.keyfile {
            buffer.extend(Sha256::digest(keyfile))
        }

        ComposedKey(Sha256::digest(&buffer).iter().cloned().collect())
    }
}

#[derive(Debug)]
/// Hashed combined input credentials used as KDF input
pub struct ComposedKey(Vec<u8>);

impl ComposedKey {
    /// Generate a master key used to derive all other keys
    pub fn master_key(
        &self,
        kdf_options: &binary::KdfParams,
    ) -> Result<MasterKey, KeyGenerationError> {
        match kdf_options {
            binary::KdfParams::Argon2 {
                memory_bytes,
                version,
                iterations,
                lanes,
                salt,
            } => {
                let config = argon2::Config {
                    variant: argon2::Variant::Argon2d,
                    version: argon2::Version::from_u32(*version)
                        .map_err(|e| KeyGenerationError::KeyGeneration(e.to_string()))?,
                    lanes: *lanes,
                    mem_cost: (memory_bytes / 1024) as u32,
                    thread_mode: argon2::ThreadMode::Parallel,
                    time_cost: *iterations as u32,
                    ..Default::default()
                };
                let hash = argon2::hash_raw(&self.0, salt, &config)
                    .map_err(|e| KeyGenerationError::KeyGeneration(e.to_string()))?;

                Ok(MasterKey(hash))
            }
            binary::KdfParams::Aes { rounds, salt } => {
                let mut cipher: Ecb<Aes256, ZeroPadding> =
                    Ecb::new_var(&salt, Default::default()).unwrap();
                let chunked: Vec<GenericArray<u8, _>> = self
                    .0
                    .chunks_exact(16)
                    .map(|chunk| GenericArray::from_slice(chunk).clone())
                    .collect();
                let mut blocks = [chunked[0], chunked[1]];
                for _ in 0..*rounds {
                    cipher.encrypt_blocks(&mut blocks);
                }
                let mut transformed_hasher = Sha256::new();
                transformed_hasher.input(blocks[0]);
                transformed_hasher.input(blocks[1]);
                let transformed = transformed_hasher.result();

                Ok(MasterKey(transformed.as_slice().to_vec()))
            }
            _ => Ok(MasterKey(Vec::new())),
        }
    }
}

/// Master key - this is generated from the user's composite key and is used to generate all other keys
#[derive(Debug)]
pub struct MasterKey(Vec<u8>);

impl MasterKey {
    /// Obtain a key to use for data integrity checks
    pub(crate) fn hmac_key(&self, seed: &[u8]) -> HmacKey {
        let mut data_to_hash = Vec::new();
        data_to_hash.extend(seed.iter());
        data_to_hash.extend(self.0.iter());
        data_to_hash.push(1);

        HmacKey(Sha512::digest(&data_to_hash).iter().cloned().collect())
    }

    /// Obtain a key to initialise a cipher
    pub(crate) fn cipher_key(&self, seed: &[u8]) -> CipherKey {
        let mut data_to_hash = Vec::new();
        data_to_hash.extend(seed.iter());
        data_to_hash.extend(self.0.iter());

        CipherKey(Sha256::digest(&data_to_hash).iter().cloned().collect())
    }
}

/// Used to initialise the encryption/decryption cipher
pub(crate) struct CipherKey(pub(crate) Vec<u8>);

/// Base key for all HMAC data integrity checks
pub(crate) struct HmacKey(Vec<u8>);

impl HmacKey {
    /// Obtain a key to verify a single block
    pub(crate) fn block_key(&self, block_idx: u64) -> HmacBlockKey {
        let mut block_key_hash = Sha512::new();
        block_key_hash.input(&block_idx.to_le_bytes());
        block_key_hash.input(&*self.0);
        HmacBlockKey(block_idx, block_key_hash.result().iter().cloned().collect())
    }
}

/// Key to perform data integrity checks on a specific block
pub(crate) struct HmacBlockKey(u64, Vec<u8>);

impl HmacBlockKey {
    /// Verify that a block in the data section is valid
    pub(crate) fn verify_data_block(&self, hmac: &[u8], data: &[u8]) -> bool {
        let mut calc_hmac = HmacSha256::new_varkey(&self.1).unwrap();
        calc_hmac.input(&self.0.to_le_bytes());
        calc_hmac.input(&(data.len() as u32).to_le_bytes());
        calc_hmac.input(data);
        match calc_hmac.verify(hmac) {
            Ok(_) => true,
            Err(_) => false,
        }
    }

    /// Calculate a HMAC for a block in the data section
    pub(crate) fn calculate_data_hmac(
        &self,
        data: &[u8],
    ) -> MacResult<<HmacSha256 as Mac>::OutputSize> {
        let mut calc_hmac = HmacSha256::new_varkey(&self.1).unwrap();
        calc_hmac.input(&self.0.to_le_bytes());
        calc_hmac.input(&(data.len() as u32).to_le_bytes());
        calc_hmac.input(data);
        calc_hmac.result()
    }

    /// Calculate a HMAC for a block in the header section
    pub(crate) fn calculate_header_hmac(
        &self,
        data: &[u8],
    ) -> MacResult<<HmacSha256 as Mac>::OutputSize> {
        let mut calc_hmac = HmacSha256::new_varkey(&self.1).unwrap();
        calc_hmac.input(data);
        calc_hmac.result()
    }

    /// Verify that the header block is valid
    pub(crate) fn verify_header_block(&self, hmac: &[u8], data: &[u8]) -> bool {
        let mut calc_hmac = HmacSha256::new_varkey(&self.1).unwrap();
        calc_hmac.input(data);
        match calc_hmac.verify(hmac) {
            Ok(_) => true,
            Err(_) => false,
        }
    }
}

/// Confirm the hash of a given block of data for data corruption detection
pub(crate) fn verify_sha256(data: &[u8], expected_sha: &[u8]) -> bool {
    expected_sha == &*Sha256::digest(&data)
}

pub(crate) fn sha256(data: &[u8]) -> Vec<u8> {
    Sha256::digest(data).as_slice().to_vec()
}

#[derive(Debug, Error)]
/// Errors encountered generating crypto keys
pub enum KeyGenerationError {
    /// Unexpected error when generating a key
    #[error("Could not generate key: {0}")]
    KeyGeneration(String),
    /// KDF Options are not supported by this library
    #[error("Generation for KDF Options: {0:?} not implemented")]
    UnimplementedKdfOptions(binary::KdfParams),
}