use c2_chacha::stream_cipher::{NewStreamCipher, SyncStreamCipher, SyncStreamCipherSeek};
use c2_chacha::{ChaCha12, ChaCha20};
use std::convert::TryInto;
use std::io::{Read, Result, Write};

pub struct ChaChaReader<'a>(ChaCha20, &'a mut Read);

impl<'a> ChaChaReader<'a> {
    pub fn new(key: &[u8], nonce: &[u8], source: &'a mut Read) -> ChaChaReader<'a> {
        ChaChaReader(ChaCha20::new_var(key, nonce).unwrap(), source)
    }
}

impl<'a> Read for ChaChaReader<'a> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let red = self.1.read(buf)?;
        self.0.apply_keystream(buf);
        Ok(red)
    }
}

pub struct ChaChaWriter<'a>(ChaCha20, &'a mut Write);

impl<'a> ChaChaWriter<'a> {
    pub fn new(key: &[u8], nonce: &[u8], sink: &'a mut Write) -> ChaChaWriter<'a> {
        ChaChaWriter(ChaCha20::new_var(key, nonce).unwrap(), sink)
    }
}

impl<'a> Write for ChaChaWriter<'a> {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        let mut cipher_text = [0u8; 256];
        let mut written = 0usize;
        for chunk in buf.chunks(cipher_text.len()) {
            cipher_text[0..chunk.len()].copy_from_slice(&chunk);
            self.0.apply_keystream(&mut cipher_text[0..chunk.len()]);
            written += self.1.write(&cipher_text[0..chunk.len()])?;
        }
        Ok(written)
    }

    fn flush(&mut self) -> Result<()> {
        self.1.flush()
    }
}