use cryptopals::{prelude::*, bvec, bvec64, english::{Bhattacharyya, self}};

fn main() {
    let s1 = b"this is a test".to_vec();
    let s2 = b"wokka wokka!!!".to_vec();
    assert_eq!(s1.hamming(&s2).unwrap(), 37);

    // infer key size
    let input = bvec64!("input.txt");

    let mut keysizes = (2..=40).map(|possible_keysize|  {
        let pk=possible_keysize;

        let bl0 = input[0*pk..1*pk].to_vec();
        let bl1 = input[1*pk..2*pk].to_vec();
        let bl2 = input[2*pk..3*pk].to_vec();
        let bl3 = input[3*pk..4*pk].to_vec();
        let bl4 = input[4*pk..5*pk].to_vec();

        let score = (
            bl0.hamming(&bl1).unwrap() +
            bl1.hamming(&bl2).unwrap() +
            bl2.hamming(&bl3).unwrap() +
            bl3.hamming(&bl4).unwrap()
        ) * 100000 / (4 * pk as u64) ;
        (pk, score)
    }).collect::<Vec<_>>();
    keysizes.sort_by_key(|(_, score)| *score);
    keysizes.drain(4..);
    dbg!(&keysizes);

    // use bhattacharyya distance to guess each key
    let sample = include_bytes!("sample.txt");
    let ideal_distribution = Bhattacharyya::compute(sample);

    let possible_keys: Vec<Vec<u8>> = keysizes.iter().map(|(keysize, _)|  {
        let key = (0..*keysize).map(|key_i| {
            let mut ciphertext_chunk=vec![];
            for i in (key_i..input.len()).step_by(*keysize) {
                ciphertext_chunk.push(input[i]);
            }

            let mut possible_bytes = (u8::MIN..=u8::MAX).map(|key_byte| {
                let plaintext = ciphertext_chunk.xor_repeating(&vec![key_byte]).unwrap();
                let score = Bhattacharyya::compute(&plaintext).score(&ideal_distribution);

                (key_byte, (score * 100000.0) as i32)
            }).collect::<Vec<_>>();

            possible_bytes.sort_by_key(|(_, sc)| -sc);
            possible_bytes.drain(1..);
            possible_bytes[0].0  // TODO: Use multiple possibilities
        }).collect::<Vec<_>>();
        key
    }).collect::<Vec<_>>();

    let mut outcomes = possible_keys.iter().map(|key| {
        let output = input.xor_repeating(&key).unwrap();
        let score = english::score(&output);
        (key, output, score)
    }).collect::<Vec<_>>();

    outcomes.sort_by_key(|(_, _, sc)| -sc);
    dbg!(outcomes[0].0.to_hex());
    dbg!(outcomes[0].1.to_text().unwrap());

    println!("pass!")
}