use crate::blake2b::Blake2b;

use super::Argon2Params;
use super::compress::{BLOCK_SIZE, BLOCK_WORDS, Block};
use super::kdf::{ARGON2_VERSION, ARGON2ID_TYPE};

/// Compute the Argon2 initial hash H_0.
///
/// For our password KDF, the optional secret and associated-data
/// inputs are both empty, but their zero lengths are still encoded
/// because they are part of the Argon2 parameter structure.
pub(super) fn initial_hash(
    password: &[u8],
    salt: &[u8],
    params: Argon2Params,
    memory_kib: u32,
    output_len: usize,
) -> [u8; 64] {
    let mut input = Vec::with_capacity(6 * 4 + 4 + password.len() + 4 + salt.len() + 4 + 4);

    // p
    append_u32(&mut input, params.parallelism);

    // T
    append_u32(&mut input, output_len as u32);

    // m'
    append_u32(&mut input, memory_kib);

    // t
    append_u32(&mut input, params.time_cost);

    // version
    append_u32(&mut input, ARGON2_VERSION);

    // type = Argon2id
    append_u32(&mut input, ARGON2ID_TYPE);

    // |P| and P
    append_u32(&mut input, password.len() as u32);
    input.extend_from_slice(password);

    // |S| and S
    append_u32(&mut input, salt.len() as u32);
    input.extend_from_slice(salt);

    // |K| = 0
    append_u32(&mut input, 0);

    // |X| = 0
    append_u32(&mut input, 0);

    let digest = Blake2b::hash(&input);

    let mut output = [0u8; 64];
    output.copy_from_slice(&digest);

    output
}

/// Compute one of the two initial memory blocks.
pub(super) fn initial_block(h0: &[u8; 64], index: usize, lane: usize) -> Block {
    let mut input = Vec::with_capacity(72);

    input.extend_from_slice(h0);

    append_u32(&mut input, index as u32);

    append_u32(&mut input, lane as u32);

    hash_long_block(&input)
}

/// Argon2 variable-length hash H'.
///
/// This is the H' construction specified by Argon2.
/// It ultimately uses the project's BLAKE2b implementation.
pub(super) fn hash_long(input: &[u8], output_len: usize) -> Vec<u8> {
    assert!(output_len >= 1);

    let mut prefixed = Vec::with_capacity(4 + input.len());

    append_u32(&mut prefixed, output_len as u32);

    prefixed.extend_from_slice(input);

    /*
     * For T <= 64:
     *
     * H'(T || A) = BLAKE2b-T(T || A)
     */
    if output_len <= 64 {
        return Blake2b::hash_with_output_length(&prefixed, output_len)
            .expect("output length <= 64");
    }

    /*
     * For T > 64:
     *
     * r = ceil(T / 32) - 2
     *
     * V_1 = H^64(T || A)
     * V_i = H^64(V_{i-1})
     * final = H^(T - 32*r)(V_r)
     */
    let r = output_len.div_ceil(32) - 2;

    let mut output = Vec::with_capacity(output_len);

    let mut v = Blake2b::hash(&prefixed);

    output.extend_from_slice(&v[..32]);

    for _ in 1..r {
        v = Blake2b::hash(&v);
        output.extend_from_slice(&v[..32]);
    }

    let final_len = output_len - 32 * r;

    v = Blake2b::hash_with_output_length(&v, final_len).expect("final H' output length is valid");

    output.extend_from_slice(&v);

    debug_assert_eq!(output.len(), output_len);

    output
}

/// Compute H'(input, 1024).
fn hash_long_block(input: &[u8]) -> Block {
    let digest = hash_long(input, BLOCK_SIZE);

    let mut words = [0u64; BLOCK_WORDS];

    for (i, word) in words.iter_mut().enumerate() {
        let start = i * 8;

        *word = u64::from_le_bytes(
            digest[start..start + 8]
                .try_into()
                .expect("Argon2 block is 1024 bytes"),
        );
    }

    Block(words)
}

/// Serialize an Argon2 block in little-endian order.
pub(super) fn block_to_bytes(block: &Block) -> Vec<u8> {
    let mut output = Vec::with_capacity(BLOCK_SIZE);

    for word in block.0 {
        output.extend_from_slice(&word.to_le_bytes());
    }

    output
}

/// Append a little-endian u32.
#[inline]
fn append_u32(output: &mut Vec<u8>, value: u32) {
    output.extend_from_slice(&value.to_le_bytes());
}