use super::compress::{Block, compress_block};
use super::hash::{block_to_bytes, hash_long, initial_block, initial_hash};
use super::kdf::ARGON2ID_TYPE;
use super::{Argon2Params, KdfError};

/// Core Argon2id function.
///
/// This implements Argon2id v1.3 using BLAKE2b for H and H'.
pub(super) fn argon2id(
    password: &[u8],
    salt: &[u8],
    params: Argon2Params,
    output_len: usize,
) -> Result<Vec<u8>, KdfError> {
    if params.time_cost == 0 {
        return Err(KdfError::InvalidTimeCost);
    }

    if params.parallelism == 0 {
        return Err(KdfError::InvalidParallelism);
    }

    let parallelism = params.parallelism as usize;
    let requested_memory = params.memory_kib as usize;

    /*
     * Argon2 requires at least 8 KiB per lane.
     */
    if requested_memory < 8 * parallelism {
        return Err(KdfError::InvalidMemoryCost);
    }

    /*
     * The actual memory size is rounded down to a multiple of 4 * parallelism.
     */
    let memory_kib = 4 * parallelism * (requested_memory / (4 * parallelism));

    /*
     * Number of blocks in each lane.
     */
    let lane_length = memory_kib / parallelism;

    /*
     * Four segments per lane.
     */
    let segment_length = lane_length / 4;

    /*
     * Initial hash H_0.
     */
    let h0 = initial_hash(password, salt, params, memory_kib as u32, output_len);

    /*
     * Allocate the complete Argon2 memory matrix.
     *
     * Layout:
     *
     *     lane 0: [block][block][block]...
     *     lane 1: [block][block][block]...
     *     ...
     */
    let mut memory = vec![Block::ZERO; memory_kib];

    /*
     * Initialize the first two blocks of every lane:
     *
     * B[i][0] = H'(H_0 || LE32(0) || LE32(i))
     * B[i][1] = H'(H_0 || LE32(1) || LE32(i))
     */
    for lane in 0..parallelism {
        memory[lane * lane_length] = initial_block(&h0, 0, lane);

        memory[lane * lane_length + 1] = initial_block(&h0, 1, lane);
    }

    /*
     * Fill the memory matrix.
     *
     * Argon2id:
     *
     * - pass 0, slices 0 and 1 use Argon2i style data independent addressing;
     * - everything else uses Argon2d style addressing.
     */
    for pass in 0..params.time_cost as usize {
        for slice in 0..4 {
            for lane in 0..parallelism {
                process_segment(
                    &mut memory,
                    lane_length,
                    segment_length,
                    parallelism,
                    pass,
                    slice,
                    lane,
                    params.time_cost as usize,
                    memory_kib as u64,
                );
            }
        }
    }

    /*
     * Final block:
     *
     * C = B[0][q-1] XOR B[1][q-1] XOR ... XOR B[p-1][q-1]
     */
    let mut final_block = Block::ZERO;

    for lane in 0..parallelism {
        let offset = lane * lane_length + lane_length - 1;

        for (dst, src) in final_block.0.iter_mut().zip(memory[offset].0.iter()) {
            *dst ^= *src;
        }
    }

    /*
     * Final Argon2 output:
     *
     * H'(C, output_len)
     */
    let final_bytes = block_to_bytes(&final_block);

    Ok(hash_long(&final_bytes, output_len))
}

/// Process one Argon2 segment.
fn process_segment(
    memory: &mut [Block],
    lane_length: usize,
    segment_length: usize,
    parallelism: usize,
    pass: usize,
    slice: usize,
    lane: usize,
    total_time: usize,
    memory_kib: u64,
) {
    /*
     * Argon2id uses data-independent addressing only during the first two slices of the first pass.
     */
    let data_independent = pass == 0 && slice < 2;

    let mut address_block = Block::ZERO;
    let mut input_block = Block::ZERO;

    /*
     * Position within this segment.
     */
    let mut index = 0usize;

    if data_independent {
        input_block.0[0] = pass as u64;
        input_block.0[1] = lane as u64;
        input_block.0[2] = slice as u64;
        input_block.0[3] = memory_kib;
        input_block.0[4] = total_time as u64;
        input_block.0[5] = ARGON2ID_TYPE as u64;

        /*
         * In the first segment, blocks 0 and 1 already exist, so generation starts at block 2.
         */
        if pass == 0 && slice == 0 {
            input_block.0[6] += 1;

            address_block = compress_block(&input_block, &Block::ZERO, None);

            address_block = compress_block(&address_block, &Block::ZERO, None);

            index = 2;
        }
    }

    while index < segment_length {
        let current_offset = lane * lane_length + slice * segment_length + index;

        /*
         * Previous block in this lane.
         *
         * At the beginning of a lane, the previous block is the final block of the same lane.
         */
        let previous_offset = if slice == 0 && index == 0 {
            current_offset + lane_length - 1
        } else {
            current_offset - 1
        };

        /*
         * Obtain J_1/J_2.
         */
        let pseudo_random = if data_independent {
            /*
             * Generate a new address block every 128 positions.
             */
            if index.is_multiple_of(128) {
                input_block.0[6] += 1;

                address_block = compress_block(&input_block, &Block::ZERO, None);

                address_block = compress_block(&address_block, &Block::ZERO, None);
            }

            address_block.0[index % 128]
        } else {
            /*
             * Argon2d addressing uses the first word of the previous block.
             */
            memory[previous_offset].0[0]
        };

        /*
         * Select the reference block.
         */
        let reference_offset = index_alpha(
            pseudo_random,
            lane_length,
            segment_length,
            parallelism,
            pass,
            slice,
            lane,
            index,
        );

        /*
         * Later passes XOR the newly generated block with the block that was already present at this position.
         */
        let previous_block = memory[previous_offset];

        let reference_block = memory[reference_offset];

        let old_block = if pass == 0 {
            None
        } else {
            Some(&memory[current_offset])
        };

        let new_block = compress_block(&previous_block, &reference_block, old_block);

        memory[current_offset] = new_block;

        index += 1;
    }
}

/// Select the reference block for an Argon2 block.
///
/// This implements the Argon2 mapping from J_1/J_2 to the reference lane and reference position.
fn index_alpha(
    random: u64,
    lane_length: usize,
    segment_length: usize,
    parallelism: usize,
    pass: usize,
    slice: usize,
    lane: usize,
    index: usize,
) -> usize {
    /*
     * J_2 selects the reference lane.
     */
    let mut reference_lane = ((random >> 32) as usize) % parallelism;

    /*
     * During the first slice of the first pass, references must remain in the same lane.
     */
    if pass == 0 && slice == 0 {
        reference_lane = lane;
    }

    /*
     * Compute the size of the allowed reference area.
     */
    let mut reference_area_size = 3 * segment_length;

    let mut start_position = ((slice + 1) % 4) * segment_length;

    if lane == reference_lane {
        reference_area_size += index;
    }

    if pass == 0 {
        reference_area_size = slice * segment_length;

        start_position = 0;

        if slice == 0 || lane == reference_lane {
            reference_area_size += index;
        }
    }

    /*
     * The current block cannot reference itself.
     */
    if index == 0 || lane == reference_lane {
        reference_area_size -= 1;
    }

    /*
     * Non-uniform mapping:
     *
     * x = J_1^2 / 2^32
     * y = |W| * x / 2^32
     */
    let mut x = (random as u32) as u64;

    x = (x * x) >> 32;

    let y = (x * reference_area_size as u64) >> 32;

    /*
     * Map into the selected lane.
     */
    let relative_position = (start_position + reference_area_size - (y as usize + 1)) % lane_length;

    reference_lane * lane_length + relative_position
}