//! NIST CAVP test vectors for arcana.
//!
//! Vectors sourced from https://github.com/coruus/nist-testvectors
//! (unofficial mirror of NIST CSRC Cryptography Toolkit).
//!
//! Run with:
//!   cargo test --release -p arcana --test cavp

use arcana::BlockCipher;

// =====================================================================
// RSP parser helpers
// =====================================================================

fn hex_to_bytes(hex: &str) -> Vec<u8> {
    (0..hex.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap())
        .collect()
}

/// Parse an RSP file into a list of (section, key-value) groups.
/// Each group is a Vec of (key, value) pairs, with section headers
/// stored as key="[section]", value="".
struct RspParser {
    groups: Vec<Vec<(String, String)>>,
}

impl RspParser {
    fn parse(path: &str) -> Self {
        let content = std::fs::read_to_string(path).unwrap_or_else(|_| panic!("Cannot read {}", path));
        let mut groups = Vec::new();
        let mut current: Vec<(String, String)> = Vec::new();

        for line in content.lines() {
            let line = line.trim();
            if line.is_empty() || line.starts_with('#') {
                continue;
            }
            if line.starts_with('[') {
                if !current.is_empty() {
                    groups.push(current.clone());
                    current.clear();
                }
                current.push((line.to_string(), String::new()));
                continue;
            }
            if let Some(pos) = line.find('=') {
                let key = line[..pos].trim().to_string();
                let val = line[pos + 1..].trim().to_string();
                current.push((key, val));
            }
        }
        if !current.is_empty() {
            groups.push(current);
        }
        Self { groups }
    }

    /// Iterate over test vectors: each vector is a group of key=value
    /// pairs between two "Count" or "Len" entries.
    fn vectors(&self) -> Vec<std::collections::HashMap<String, String>> {
        let mut result = Vec::new();
        let mut current = std::collections::HashMap::new();
        let mut section = String::new();

        for group in &self.groups {
            for (key, val) in group {
                if key.starts_with('[') {
                    section = key.clone();
                    continue;
                }
                if key == "Count" || key == "COUNT" || (key == "Len" && !current.is_empty()) {
                    if !current.is_empty() {
                        current.insert("_section".to_string(), section.clone());
                        result.push(current.clone());
                        current.clear();
                    }
                }
                current.insert(key.clone(), val.clone());
            }
        }
        if !current.is_empty() {
            current.insert("_section".to_string(), section);
            result.push(current);
        }
        result
    }
}

fn get(v: &std::collections::HashMap<String, String>, k: &str) -> Vec<u8> {
    hex_to_bytes(v.get(k).unwrap_or(&String::new()))
}

fn get_str<'a>(v: &'a std::collections::HashMap<String, String>, k: &str) -> &'a str {
    v.get(k).map(|s| s.as_str()).unwrap_or("")
}

// =====================================================================
// SHA-1 / SHA-256 / SHA-384 / SHA-512 ShortMsg
// =====================================================================

fn cavp_sha_short_msg<H: arcana::Hasher>(path: &str, label: &str) {
    let rsp = RspParser::parse(path);
    let vecs = rsp.vectors();
    let mut count = 0u32;

    for v in &vecs {
        let len_bits: usize = v.get("Len").unwrap().parse().unwrap();
        let msg = get(v, "Msg");
        let md = get(v, "MD");

        // Skip bit-oriented lengths that aren't byte-aligned.
        if len_bits % 8 != 0 {
            continue;
        }
        let len_bytes = len_bits / 8;
        let input = &msg[..len_bytes];

        let digest = H::hash(input);
        assert_eq!(digest, md, "{} CAVP: mismatch at Len={}", label, len_bits);
        count += 1;
    }

    println!("CAVP {}: {} vectors passed", label, count);
    assert!(count > 30, "expected > 30 {} vectors", label);
}

#[test]
fn cavp_sha1_short_msg() {
    cavp_sha_short_msg::<arcana::hash::sha1::Sha1>("tests/cavp/SHA1ShortMsg.rsp", "SHA-1");
}

#[test]
fn cavp_sha256_short_msg() {
    cavp_sha_short_msg::<arcana::hash::sha256::Sha256>("tests/cavp/SHA256ShortMsg.rsp", "SHA-256");
}

#[test]
fn cavp_sha384_short_msg() {
    cavp_sha_short_msg::<arcana::hash::sha384::Sha384>("tests/cavp/SHA384ShortMsg.rsp", "SHA-384");
}

#[test]
fn cavp_sha512_short_msg() {
    cavp_sha_short_msg::<arcana::hash::sha512::Sha512>("tests/cavp/SHA512ShortMsg.rsp", "SHA-512");
}

#[test]
fn cavp_sha256_long_msg() {
    cavp_sha_short_msg::<arcana::hash::sha256::Sha256>("tests/cavp/SHA256LongMsg.rsp", "SHA-256 LongMsg");
}

// =====================================================================
// AES-CBC MMT (Multi-block Message Test)
// =====================================================================

fn cavp_aes_cbc_mmt(path: &str, key_len: usize, label: &str) {
    use arcana::cipher::aes::Aes;
    use arcana::cipher::modes::{cbc_decrypt, cbc_encrypt};

    let rsp = RspParser::parse(path);
    let vecs = rsp.vectors();
    let mut enc_count = 0u32;
    let mut dec_count = 0u32;

    for v in &vecs {
        let section = get_str(v, "_section");
        let key = get(v, "KEY");
        let iv = get(v, "IV");

        if key.len() != key_len {
            continue;
        }

        let cipher = Aes::new(&key);

        if section.contains("ENCRYPT") {
            let pt = get(v, "PLAINTEXT");
            let ct_expected = get(v, "CIPHERTEXT");
            let mut data = pt.clone();
            cbc_encrypt(&cipher, &iv, &mut data);
            assert_eq!(data, ct_expected, "{} encrypt mismatch", label);
            enc_count += 1;
        } else if section.contains("DECRYPT") {
            let ct = get(v, "CIPHERTEXT");
            let pt_expected = get(v, "PLAINTEXT");
            let mut data = ct.clone();
            cbc_decrypt(&cipher, &iv, &mut data);
            assert_eq!(data, pt_expected, "{} decrypt mismatch", label);
            dec_count += 1;
        }
    }

    println!("CAVP {}: {} encrypt + {} decrypt passed", label, enc_count, dec_count);
    assert!(enc_count > 5, "expected > 5 {} encrypt vectors", label);
    assert!(dec_count > 5, "expected > 5 {} decrypt vectors", label);
}

#[test]
fn cavp_aes128_cbc_mmt() {
    cavp_aes_cbc_mmt("tests/cavp/CBCMMT128.rsp", 16, "AES-128-CBC");
}

#[test]
fn cavp_aes256_cbc_mmt() {
    cavp_aes_cbc_mmt("tests/cavp/CBCMMT256.rsp", 32, "AES-256-CBC");
}

// =====================================================================
// AES-ECB MMT
// =====================================================================

#[test]
fn cavp_aes128_ecb_mmt() {
    use arcana::cipher::aes::Aes;
    use arcana::cipher::modes::{ecb_decrypt, ecb_encrypt};

    let rsp = RspParser::parse("tests/cavp/ECBMMT128.rsp");
    let vecs = rsp.vectors();
    let mut count = 0u32;

    for v in &vecs {
        let section = get_str(v, "_section");
        let key = get(v, "KEY");
        if key.len() != 16 {
            continue;
        }
        let cipher = Aes::new(&key);

        if section.contains("ENCRYPT") {
            let pt = get(v, "PLAINTEXT");
            let ct_expected = get(v, "CIPHERTEXT");
            let mut data = pt.clone();
            ecb_encrypt(&cipher, &mut data);
            assert_eq!(data, ct_expected, "AES-128-ECB encrypt mismatch");
            count += 1;
        } else if section.contains("DECRYPT") {
            let ct = get(v, "CIPHERTEXT");
            let pt_expected = get(v, "PLAINTEXT");
            let mut data = ct.clone();
            ecb_decrypt(&cipher, &mut data);
            assert_eq!(data, pt_expected, "AES-128-ECB decrypt mismatch");
            count += 1;
        }
    }

    println!("CAVP AES-128-ECB: {} vectors passed", count);
    assert!(count > 10, "expected > 10 AES-128-ECB vectors");
}

// =====================================================================
// HMAC (L=20 SHA-1, L=32 SHA-256, L=48 SHA-384, L=64 SHA-512)
// =====================================================================

#[test]
fn cavp_hmac() {
    use arcana::mac::ctx::{Algorithm, Mac};

    let content = std::fs::read_to_string("tests/cavp/HMAC.rsp").unwrap();
    let mut counts = [0u32; 4]; // sha1, sha256, sha384, sha512
    let mut current_l = 0usize;
    let mut kv = std::collections::HashMap::new();

    for line in content.lines() {
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        if line.starts_with("[L=") {
            current_l = line[3..line.len() - 1].parse().unwrap_or(0);
            continue;
        }
        if let Some(pos) = line.find('=') {
            let key = line[..pos].trim();
            let val = line[pos + 1..].trim();
            kv.insert(key.to_string(), val.to_string());
        }
        // When we have all fields for a vector, process it.
        if kv.contains_key("Mac") && kv.contains_key("Key") && kv.contains_key("Msg") {
            let (algo, idx) = match current_l {
                20 => (Algorithm::HmacSha1, 0),
                32 => (Algorithm::HmacSha256, 1),
                48 => (Algorithm::HmacSha384, 2),
                64 => (Algorithm::HmacSha512, 3),
                _ => {
                    kv.clear();
                    continue;
                } // L=28 (SHA-224) — skip
            };

            let key_bytes = hex_to_bytes(kv.get("Key").unwrap());
            let msg_bytes = hex_to_bytes(kv.get("Msg").unwrap());
            let expected = hex_to_bytes(kv.get("Mac").unwrap());
            let tlen: usize = kv.get("Tlen").unwrap().parse().unwrap();

            let mut m = Mac::new(algo);
            m.init(&key_bytes).unwrap();
            m.update(&msg_bytes).unwrap();
            let full_tag = m.sign_to_vec().unwrap();
            let truncated = &full_tag[..tlen];

            assert_eq!(
                truncated,
                &expected[..],
                "CAVP HMAC L={} mismatch at Count={}",
                current_l,
                kv.get("Count").unwrap_or(&"?".to_string()),
            );
            counts[idx] += 1;
            kv.clear();
        }
    }

    println!(
        "CAVP HMAC: SHA-1={}, SHA-256={}, SHA-384={}, SHA-512={}",
        counts[0], counts[1], counts[2], counts[3]
    );
    for (i, label) in ["SHA-1", "SHA-256", "SHA-384", "SHA-512"].iter().enumerate() {
        assert!(counts[i] > 50, "expected > 50 HMAC-{} vectors", label);
    }
}

// =====================================================================
// CMAC-AES-128 Generation
// =====================================================================

#[test]
fn cavp_cmac_aes128() {
    use arcana::mac::ctx::{Algorithm, Mac};

    let rsp = RspParser::parse("tests/cavp/CMACGenAES128.rsp");
    let vecs = rsp.vectors();
    let mut count = 0u32;

    for v in &vecs {
        let key = get(v, "Key");
        let msg = get(v, "Msg");
        let expected = get(v, "Mac");
        let tlen: usize = v.get("Tlen").unwrap_or(&"16".to_string()).parse().unwrap();

        if key.len() != 16 {
            continue;
        }

        let mut m = Mac::new(Algorithm::CmacAes128);
        m.init(&key).unwrap();
        // CAVP CMAC: Mlen=0 means empty message, but Msg field may
        // still be present as "00".
        let mlen: usize = v.get("Mlen").unwrap_or(&"0".to_string()).parse().unwrap();
        if mlen > 0 {
            m.update(&msg[..mlen]).unwrap();
        } else {
            m.update(&[]).unwrap();
        }
        let full_tag = m.sign_to_vec().unwrap();
        let truncated = &full_tag[..tlen];

        assert_eq!(
            truncated,
            &expected[..],
            "CAVP CMAC-AES-128 mismatch at Count={}",
            get_str(v, "Count"),
        );
        count += 1;
    }

    println!("CAVP CMAC-AES-128: {} vectors passed", count);
    assert!(count > 20, "expected > 20 CMAC-AES-128 vectors");
}

// =====================================================================
// AES-GCM (limited to first 500 vectors to keep test time reasonable)
// =====================================================================

#[test]
fn cavp_aes128_gcm() {
    use arcana::cipher::aes::Aes128;
    use arcana::cipher::modes::Gcm;

    let rsp = RspParser::parse("tests/cavp/gcmEncryptExtIV128.rsp");
    let vecs = rsp.vectors();
    let mut count = 0u32;
    let max_vectors = 500;

    for v in &vecs {
        if count >= max_vectors {
            break;
        }

        let key = get(v, "Key");
        let iv = get(v, "IV");
        let pt = get(v, "PT");
        let aad = get(v, "AAD");
        let ct_expected = get(v, "CT");
        let tag_expected = get(v, "Tag");

        // Only test 96-bit IV and 128-bit tag.
        if key.len() != 16 || iv.len() != 12 || tag_expected.len() != 16 {
            continue;
        }

        let nonce: [u8; 12] = iv.try_into().unwrap();
        let cipher = Aes128::new(&key);
        let (ct, tag) = Gcm::encrypt(&cipher, &nonce, &aad, &pt);

        assert_eq!(ct, ct_expected, "CAVP GCM encrypt CT mismatch");
        assert_eq!(tag.to_vec(), tag_expected, "CAVP GCM encrypt tag mismatch");
        count += 1;
    }

    println!("CAVP AES-128-GCM: {} vectors passed", count);
    assert!(count > 100, "expected > 100 AES-128-GCM vectors");
}
