use cryptopals::{set2_adversary::Set2Adversary, friendly::likely_ecb, prelude::ByteBased};

fn main() {
    let adversary = Set2Adversary::gen();

    // first of all, find the block size
    let block_size = {
        let mut inp = vec![];
        let l1 = adversary.oracle_12(&inp).len();
        loop {
            inp.push(b'A');
            let l2 = adversary.oracle_12(&inp).len();
            if l2 > l1 {
                break l2 - l1;
            }
        }
    };

    assert_eq!(block_size, 16);

    // make sure we're in ECB mode
    assert!(likely_ecb(&adversary.oracle_12(&[b'A'; 64])));

    let message_len = adversary.oracle_12(&[]).len();

    let mut known_message: Vec<u8> = vec![];

    // sample questions for 'ABCD EFGH' and block size 4
    // question   dictionary pattern
    // 000        000[A-Z]
    // 00         00A[A-Z]
    // 0          0AB[A-Z]
    //            ABC[A-Z]
    // 000        BCD[A-Z]
    // 00         CDE[A-Z]
    // 0          DEF[A-Z]
    //            EFG[A-Z]

    'inspect: while known_message.len() < message_len {
        let interest_ix = known_message.len();

        let interest_block_ix= interest_ix / block_size;
        let interest_block_byte_ix = interest_ix % block_size;
        let interest_shiftr1_padding = (block_size - interest_block_byte_ix - 1) % block_size; // padding it takes to put the character of interest at the end of a block

        // build shiftr
        let mut shiftr_block = vec![];
        for _ in 0..interest_shiftr1_padding { shiftr_block.push(0)}

        // build question
        // this is the beginning of the dictionary pattern (minus the range at the end)
        // if we're at the beginning of the string, use a zero as we know this
        // comes from our shiftr block
        // otherwise, use a character from the end of the string
        let mut question_block = vec![];
        for i in 0..block_size-1 {
            let source_char_ix = known_message.len() as i64 - block_size as i64 + i as i64 + 1;
            if source_char_ix < 0 { question_block.push(0); }
            else { question_block.push(known_message[source_char_ix as usize]); }
        }

        // build dictionary
        let mut dictionary: Vec<Vec<u8>> = vec![];
        for possible_byte in u8::MIN..=u8::MAX {
            let mut dict_block = question_block.clone();
            dict_block.push(possible_byte);

            let mut oracle_result = adversary.oracle_12(&dict_block);
            dictionary.push(oracle_result.drain(0..block_size).collect())
        }

        // now shift the input. we know the position of the block containing the byte of interest
        let mut tea_leaves: Vec<u8> = adversary.oracle_12(&shiftr_block);
        let output_len = tea_leaves.len();
        let block_of_interest: Vec<u8> = tea_leaves.drain(interest_block_ix * block_size..interest_block_ix * block_size + block_size).collect();

        for possible_byte in u8::MIN..=u8::MAX {
            if dictionary[possible_byte as usize] == block_of_interest {
                if (possible_byte as usize) < block_size && (interest_ix + possible_byte as usize + shiftr_block.len() == output_len) {
                    break 'inspect  // we hit padding
                }
                known_message.push(possible_byte);
                continue 'inspect
            }
        }

        panic!("we failed to break on padding");
    }

    println!("{}", known_message.to_text().unwrap())
}