use aes::Aes256;
use fpe::ff1::{FF1, FlexibleNumeralString};

#[derive(Clone)]
pub struct Scrambler {
	aes_key: [u8; 32],
	//	aes_key_fingerprint: String,
}

#[derive(Debug)]
pub enum ScramblerError {
	InvalidCharacter,
}

impl Scrambler {
	pub fn new(aes_key: [u8; 32]) -> Scrambler {
		Scrambler { aes_key }
	}

	pub fn id_to_state_key(&self, mut id: u64) -> String {
		let mut id62: Vec<u16> = vec![];
		while id > 0 {
			id62.push((id % 62) as u16);
			id = id / 62;
		}
		while id62.len() < 4 {
			id62.push(0)
		}
		let encrypted: Vec<u16> = FF1::<Aes256>::new(&self.aes_key, 62).unwrap().encrypt(&[], &FlexibleNumeralString::from(id62)).unwrap().into();
		encrypted
			.into_iter()
			.map(|v| {
				[
					'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R',
					'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't',
					'u', 'v', 'w', 'x', 'y', 'z',
				][v as usize]
			})
			.collect()
	}

	pub fn state_key_to_id(&self, key: &str) -> Result<u64, ScramblerError> {
		let encrypted: Vec<u16> = key
			.chars()
			.map(|char| match char {
				'0'..='9' => Ok(char as u16 - '0' as u16),
				'A'..='Z' => Ok(char as u16 - 'A' as u16 + 10),
				'a'..='z' => Ok(char as u16 - 'a' as u16 + 36),
				_ => return Err(ScramblerError::InvalidCharacter),
			})
			.collect::<Result<_, ScramblerError>>()?;
		let decrypted: Vec<u16> = FF1::<Aes256>::new(&self.aes_key, 62).unwrap().decrypt(&[], &FlexibleNumeralString::from(encrypted)).unwrap().into();
		Ok(decrypted.iter().enumerate().fold(0 as u64, |acc, (i, v)| acc + (62 as u64).pow(i as u32) * *v as u64))
	}
}

#[cfg(test)]
mod tests {
	use rand::{Rng, SeedableRng};

	use super::*;

	fn roundtrip(fpe: &Scrambler, original_id: u64) {
		let state_key = fpe.id_to_state_key(original_id);
		let id = fpe.state_key_to_id(&state_key).unwrap();
		assert_eq!(original_id, id);
	}

	#[test]
	fn test_id_to_state_key() {
		let fpe = Scrambler::new([0u8; 32]);
		roundtrip(&fpe, u64::MIN);
		roundtrip(&fpe, u64::MAX);
		let mut rng = rand::rngs::StdRng::seed_from_u64(42);
		for _ in 0..10000 {
			roundtrip(&fpe, rng.random());
		}
	}
}
