use argon2::{
    Argon2,
    password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString, rand_core::OsRng},
};
use bytes::Bytes;
use sha2::{Digest, Sha256};
use std::error::Error;

/// generates 6 recovery codes and their Argon2 hashes
pub fn generate_codes() -> Result<(Bytes, String), Box<dyn Error>> {
    // total len == 55
    let mut recovery_codes = String::with_capacity(55);

    let mut recovery_code_hashes_builder =
        flexbuffers::Builder::new(&flexbuffers::BuilderOptions::SHARE_NONE);
    let mut recovery_code_hashes_builder_vec = recovery_code_hashes_builder.start_vector();

    let argon2 = Argon2::default();
    let salt = SaltString::generate(&mut OsRng);

    for i in 0..3 {
        // len == 22
        let random_code = SaltString::generate(&mut OsRng);

        // len == 11
        let mut code1 = String::with_capacity(11);
        let mut code2 = String::with_capacity(11);

        for (j, c) in random_code.to_string().chars().enumerate() {
            if j < 11 {
                recovery_codes.push(c);
                code1.push(c);
            } else {
                if i == 2 {
                    break;
                }

                recovery_codes.push(c);
                code2.push(c);
            }
        }

        let mut hasher1 = Sha256::new();
        hasher1.update(code1.as_bytes());
        let hash_code1 = hasher1.finalize().to_vec();

        let recovery_code1_hash = match argon2.hash_password(code1.as_bytes(), &salt) {
            Ok(v) => v,
            Err(err) => return Err(err.to_string().into()),
        }
        .to_string();
        let recovery_hash_code1_hash = match argon2.hash_password(&hash_code1, &salt) {
            Ok(v) => v,
            Err(err) => return Err(err.to_string().into()),
        }
        .to_string();

        let mut code_pair1 = recovery_code_hashes_builder_vec.start_vector();

        code_pair1.push(recovery_hash_code1_hash.as_str());
        code_pair1.push(recovery_code1_hash.as_str());

        code_pair1.end_vector();

        if i != 2 {
            let mut hasher2 = Sha256::new();
            hasher2.update(code2.as_bytes());
            let hash_code2 = hasher2.finalize().to_vec();

            let recovery_code2_hash = match argon2.hash_password(code2.as_bytes(), &salt) {
                Ok(v) => v,
                Err(err) => return Err(err.to_string().into()),
            }
            .to_string();
            let recovery_hash_code2_hash = match argon2.hash_password(&hash_code2, &salt) {
                Ok(v) => v,
                Err(err) => return Err(err.to_string().into()),
            }
            .to_string();

            let mut code_pair2 = recovery_code_hashes_builder_vec.start_vector();

            code_pair2.push(recovery_hash_code2_hash.as_str());
            code_pair2.push(recovery_code2_hash.as_str());

            code_pair2.end_vector();
        }
    }

    recovery_code_hashes_builder_vec.end_vector();

    Ok((
        Bytes::copy_from_slice(recovery_code_hashes_builder.view()),
        recovery_codes,
    ))
}

pub fn check_code(
    hashed_recovery_code: &[u8],
    stored_codes: &[u8],
) -> Result<bool, Box<dyn Error>> {
    let recovery_codes_root = flexbuffers::Reader::get_root(stored_codes)?;

    let argon2 = Argon2::default();
    let mut is_valid_recovery_code = false;

    for code_pair in &recovery_codes_root.as_vector() {
        if let Ok(parsed_hash) = PasswordHash::new(code_pair.as_vector().idx(0).as_str())
            && let Ok(()) = argon2.verify_password(hashed_recovery_code, &parsed_hash)
        {
            is_valid_recovery_code = true;
            break;
        }
    }

    Ok(is_valid_recovery_code)
}

pub fn consume_code(
    recovery_code: &[u8],
    stored_codes: &[u8],
) -> Result<(bool, Bytes), Box<dyn Error>> {
    let recovery_codes_root = flexbuffers::Reader::get_root(stored_codes)?;

    let mut write_recovery_codes =
        flexbuffers::Builder::new(&flexbuffers::BuilderOptions::SHARE_NONE);
    let mut write_recovery_codes_vec = write_recovery_codes.start_vector();

    let argon2 = Argon2::default();
    let mut is_valid_recovery_code = false;

    for code_pair in &recovery_codes_root.as_vector() {
        let code_pair_vec = code_pair.as_vector();

        if is_valid_recovery_code {
            let mut pair = write_recovery_codes_vec.start_vector();

            pair.push(code_pair_vec.idx(0).as_str());
            pair.push(code_pair_vec.idx(1).as_str());

            pair.end_vector();
        } else if let Ok(parsed_hash) = PasswordHash::new(code_pair_vec.idx(1).as_str()) {
            if argon2.verify_password(recovery_code, &parsed_hash).is_ok() {
                is_valid_recovery_code = true;
            } else {
                let mut pair = write_recovery_codes_vec.start_vector();

                pair.push(code_pair_vec.idx(0).as_str());
                pair.push(code_pair_vec.idx(1).as_str());

                pair.end_vector();
            }
        }
    }

    write_recovery_codes_vec.end_vector();

    Ok((
        is_valid_recovery_code,
        Bytes::copy_from_slice(write_recovery_codes.view()),
    ))
}
