use aes::cipher::generic_array::GenericArray;
use aes::cipher::BlockEncrypt;
use std::io::{Read, Result as IoResult};

pub struct AesDecoder<Io: Read, Aes: BlockEncrypt> {
    io: Io,
    aes: Aes,

    counter: u128,
    block: [u8; 16],
    lower: usize,
    upper: usize,
}

impl<Io: Read, Aes: BlockEncrypt> AesDecoder<Io, Aes> {
    pub fn new(io: Io, aes: Aes) -> IoResult<Self> {
        let mut decoder = Self {
            io,
            aes,

            counter: 0,
            block: [0; 16],
            lower: 0,
            upper: 0,
        };

        decoder.update_block()?;

        Ok(decoder)
    }

    #[inline]
    fn update_block(&mut self) -> IoResult<()> {
        self.upper = self.io.read(&mut self.block)?;
        self.lower = 0;

        self.counter += 1;
        let mut mask = self.counter.to_le_bytes();

        self.aes
            .encrypt_block(GenericArray::from_mut_slice(&mut mask));

        for (b, m) in self.block.iter_mut().zip(mask) {
            *b ^= m
        }

        Ok(())
    }
}

impl<Io: Read, Aes: BlockEncrypt> Read for AesDecoder<Io, Aes> {
    fn read(&mut self, mut buf: &mut [u8]) -> IoResult<usize> {
        let mut bytes = 0;
        while !buf.is_empty() && self.lower != self.upper {
            for (to, fr) in buf.iter_mut().zip(&self.block[self.lower..self.upper]) {
                *to = *fr
            }

            let consumed = buf.len().min(self.upper - self.lower);
            buf = &mut buf[consumed..];
            self.lower += consumed;
            bytes += consumed;

            if self.lower == 16 {
                self.update_block()?;
            }
        }
        Ok(bytes)
    }
}