use rayon::prelude::*;
use std::collections::HashMap;

use crate::tea3::Tea3;

/// period brute force with hashmap
pub fn period_bruteforce(key: Vec<u8>, iv: Vec<u8>) -> u64 {
    let mut cipher = Tea3::new(key, iv);
    cipher.init();

    let mut seen = HashMap::new();
    let mut step: u64 = 0;

    loop {
        let state = (
            cipher.key_register().to_vec(),
            cipher.state_register().to_vec(),
        );

        if let Some(prev) = seen.get(&state) {
            return step - prev;
        }

        seen.insert(state, step);

        cipher.step();
        step += 1;
    }
}

/// period brute force with Floyd tortoise and hare algorithm
pub fn period_floyd(key: Vec<u8>, iv: Vec<u8>) -> u64 {
    let mut tortoise = Tea3::new(key.clone(), iv.clone());
    let mut hare = Tea3::new(key, iv);

    tortoise.init();
    hare.init();

    tortoise.step();
    hare.step();
    hare.step();

    while (tortoise.key_register(), tortoise.state_register())
        != (hare.key_register(), hare.state_register())
    {
        tortoise.step();
        hare.step();
        hare.step();
    }

    let mut lambda = 1;
    hare.step();

    while (tortoise.key_register(), tortoise.state_register())
        != (hare.key_register(), hare.state_register())
    {
        hare.step();
        lambda += 1;
    }

    lambda
}

// paper method, decomposition key register
// (Observations on TETRA Encryption Algorithm TEA-3)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Reg5([u8; 5]);

impl Reg5 {
    fn from_u40(x: u64) -> Self {
        assert!(x < (1u64 << 40), "value must fit in 40 bits");

        // Internal order: byte 0 is the first register byte.
        Self([
            ((x >> 32) & 0xFF) as u8,
            ((x >> 24) & 0xFF) as u8,
            ((x >> 16) & 0xFF) as u8,
            ((x >> 8) & 0xFF) as u8,
            (x & 0xFF) as u8,
        ])
    }

    fn step_h(&mut self) -> u8 {
        let out = self.0[0];
        let feedback = out ^ Tea3::p(self.0[2]);

        self.0.copy_within(1..5, 0);
        self.0[4] = feedback;

        out
    }

    fn step_g(&mut self, h_out: u8) -> u8 {
        let out = self.0[0];
        let feedback = out ^ h_out;

        self.0.copy_within(1..5, 0);
        self.0[4] = feedback;

        out
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Decomposed {
    g: Reg5,
    h: Reg5,
}

impl Decomposed {
    fn from_key(key: &[u8]) -> Self {
        Self {
            g: Reg5([key[0], key[1], key[2], key[3], key[4]]),
            h: Reg5([
                key[0] ^ key[5],
                key[1] ^ key[6],
                key[2] ^ key[7],
                key[3] ^ key[8],
                key[4] ^ key[9],
            ]),
        }
    }

    fn step(&mut self) -> u8 {
        let h_out = self.h.step_h();
        self.g.step_g(h_out)
    }
}

fn period_h(start: Reg5) -> u64 {
    let mut tortoise = start;
    let mut hare = start;

    tortoise.step_h();
    hare.step_h();
    hare.step_h();

    while tortoise != hare {
        tortoise.step_h();
        hare.step_h();
        hare.step_h();
    }

    let mut lambda = 1;
    hare.step_h();

    while tortoise != hare {
        hare.step_h();
        lambda += 1;
    }

    lambda
}

/// paper method, decomposition key register
/// (Observations on TETRA Encryption Algorithm TEA-3)
/// key 5 bytes, (10 bytes but first 5 bytes are padding with 0)
pub fn period_paper(key: Vec<u8>) -> u64 {
    let start = Decomposed::from_key(&[0, 0, 0, 0, 0, key[0], key[1], key[2], key[3], key[4]]);
    period_h(start.h) // p
}

/// Exhaustively scan all 2^40 initial values of h and returns:
/// (longest period found, 5-bytes initial value)
pub fn longest_period() -> (u64, [u8; 5]) {
    let mut best_period = 0u64;
    let mut best_seed = [0u8; 5];

    for x in 0u64..(1u64 << 40) {
        let seed = Reg5::from_u40(x);
        let p = period_h(seed);

        if p > best_period {
            best_period = p;
            best_seed = seed.0;
        }
    }

    (best_period, best_seed)
}

pub fn longest_period_parallel() -> (u64, [u8; 5]) {
    let limit: usize = 1usize << 40;

    (0usize..limit)
        .into_par_iter()
        .map(|x| {
            let seed = Reg5::from_u40(x as u64);
            let p = period_h(seed);
            (p, seed.0)
        })
        .reduce(|| (0u64, [0u8; 5]), |a, b| if a.0 >= b.0 { a } else { b })
}