pub use private::SecretStream;

mod private {
    use crate::{
        crypto::{Error, Result, SymKey},
        Block, BlockMeta, Decompose, MetaAccess, Sectored, TryCompose,
    };
    use std::io::{self, Read, Seek, SeekFrom, Write};

    // A stream which encrypts all data written to it and decrypts all data read from it.
    pub struct SecretStream<T> {
        inner: T,
        // The sector size of the inner stream. Reads and writes are only executed using buffers of
        // this size.
        inner_sect_sz: usize,
        // The sector size of this stream. Reads and writes are only accepted for buffers of this size.
        sect_sz: usize,
        key: SymKey,
        /// Buffer for ciphertext.
        ct_buf: Vec<u8>,
        /// Buffer for plaintext.
        pt_buf: Vec<u8>,
    }

    impl<T> SecretStream<T> {
        /// Given an offset into this stream, produces the corresponding offset into the inner stream.
        fn inner_offset(&self, outer_offset: u64) -> u64 {
            let sect_sz = self.sect_sz as u64;
            let inner_sect_sz = self.inner_sect_sz as u64;
            // We return the offset into the current sector, plus the size of all previous sectors.
            outer_offset % sect_sz + outer_offset / sect_sz * inner_sect_sz
        }

        /// Given an offset into the inner stream, returns the corresponding offset into this stream.
        fn outer_offset(&self, inner_offset: u64) -> u64 {
            let sect_sz = self.sect_sz as u64;
            let inner_sect_sz = self.inner_sect_sz as u64;
            inner_offset % inner_sect_sz + inner_offset / inner_sect_sz * sect_sz
        }
    }

    impl SecretStream<()> {
        pub fn new(key: SymKey) -> SecretStream<()> {
            SecretStream {
                inner: (),
                inner_sect_sz: 0,
                sect_sz: 0,
                key,
                ct_buf: Vec::new(),
                pt_buf: Vec::new(),
            }
        }
    }

    impl<T> Decompose<T> for SecretStream<T> {
        fn into_inner(self) -> T {
            self.inner
        }
    }

    impl<T, U: Sectored> TryCompose<U, SecretStream<U>> for SecretStream<T> {
        type Error = Error;
        fn try_compose(mut self, inner: U) -> Result<SecretStream<U>> {
            let inner_sect_sz = inner.sector_sz();
            let expansion_sz = self.key.expansion_sz();
            let sect_sz = inner_sect_sz - expansion_sz;
            let block_sz = self.key.block_size();
            if 0 != sect_sz % block_sz {
                return Err(Error::IndivisibleSize {
                    divisor: block_sz,
                    actual: sect_sz,
                });
            }
            self.pt_buf.resize(inner_sect_sz, 0);
            self.pt_buf.resize(inner_sect_sz + block_sz, 0);
            Ok(SecretStream {
                inner,
                inner_sect_sz,
                sect_sz: inner_sect_sz - expansion_sz,
                key: self.key,
                ct_buf: self.ct_buf,
                pt_buf: self.pt_buf,
            })
        }
    }

    impl<T> Sectored for SecretStream<T> {
        fn sector_sz(&self) -> usize {
            self.sect_sz
        }
    }

    impl<T: Write> Write for SecretStream<T> {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.assert_sector_sz(buf.len())?;

            self.ct_buf.resize(self.inner_sect_sz, 0);
            let mut encrypter = self.key.to_encrypter()?;
            let mut count = encrypter.update(buf, &mut self.ct_buf)?;
            count += encrypter.finalize(&mut self.ct_buf[count..])?;
            self.ct_buf.truncate(count);

            self.inner.write_all(&self.ct_buf).map(|_| buf.len())
        }

        fn flush(&mut self) -> io::Result<()> {
            self.inner.flush()
        }
    }

    impl<T: Read> Read for SecretStream<T> {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            self.assert_sector_sz(buf.len())?;

            self.ct_buf.resize(self.inner_sect_sz, 0);
            match self.inner.read_exact(&mut self.ct_buf) {
                Ok(_) => (),
                Err(err) => {
                    if err.kind() == io::ErrorKind::UnexpectedEof {
                        return Ok(0);
                    } else {
                        return Err(err);
                    }
                }
            }

            self.pt_buf
                .resize(self.inner_sect_sz + self.key.block_size(), 0);
            let mut decrypter = self.key.to_decrypter()?;
            let mut count = decrypter.update(&self.ct_buf, &mut self.pt_buf)?;
            count += decrypter.finalize(&mut self.pt_buf[count..])?;
            self.pt_buf.truncate(count);

            buf.copy_from_slice(&self.pt_buf);
            Ok(buf.len())
        }
    }

    impl<T: Seek> Seek for SecretStream<T> {
        fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
            let outer_offset = match pos {
                SeekFrom::Start(offset) => offset,
                SeekFrom::Current(offset) => {
                    let inner_offset = self.inner.stream_position()?;
                    let outer_offset = self.outer_offset(inner_offset);
                    if offset >= 0 {
                        outer_offset + offset as u64
                    } else {
                        outer_offset - (-offset as u64)
                    }
                }
                SeekFrom::End(_) => {
                    // We can support this once stream_len is stabilized:
                    // https://github.com/rust-lang/rust/issues/59359
                    return Err(io::Error::new(
                        io::ErrorKind::Unsupported,
                        "seeking from the end of the stream is not supported",
                    ));
                }
            };
            let inner_offset = self.inner_offset(outer_offset);
            self.inner.seek(SeekFrom::Start(inner_offset))?;
            Ok(outer_offset)
        }
    }

    impl<T: AsRef<BlockMeta>> AsRef<BlockMeta> for SecretStream<T> {
        fn as_ref(&self) -> &BlockMeta {
            self.inner.as_ref()
        }
    }

    impl<T: AsMut<BlockMeta>> AsMut<BlockMeta> for SecretStream<T> {
        fn as_mut(&mut self) -> &mut BlockMeta {
            self.inner.as_mut()
        }
    }

    impl<T: MetaAccess> MetaAccess for SecretStream<T> {}

    impl<T: Block> Block for SecretStream<T> {
        fn flush_meta(&mut self) -> crate::Result<()> {
            self.inner.flush_meta()
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::{Read, Seek, SeekFrom, Write};

    use crate::{
        crypto::{SymKey, SymKeyKind},
        test_helpers::{Randomizer, SectoredCursor},
        Sectored, TryCompose, SECTOR_SZ_DEFAULT,
    };

    use super::*;

    fn secret_stream_sequential_test_case(key: SymKey, inner_sect_sz: usize, sect_ct: usize) {
        let mut stream = SecretStream::new(key)
            .try_compose(SectoredCursor::new(
                vec![0u8; inner_sect_sz * sect_ct],
                inner_sect_sz,
            ))
            .expect("compose failed");
        let sector_sz = stream.sector_sz();
        for k in 0..sect_ct {
            let sector = vec![k as u8; sector_sz];
            stream.write(&sector).expect("write failed");
        }
        stream.seek(SeekFrom::Start(0)).expect("seek failed");
        for k in 0..sect_ct {
            let expected = vec![k as u8; sector_sz];
            let mut actual = vec![0u8; sector_sz];
            stream.read(&mut actual).expect("read failed");
            assert_eq!(expected, actual);
        }
    }

    fn secret_stream_sequential_test_suite(kind: SymKeyKind) {
        let key = SymKey::generate(kind).expect("key generation failed");
        secret_stream_sequential_test_case(key.clone(), SECTOR_SZ_DEFAULT, 16);
    }

    #[test]
    fn secret_stream_encrypt_decrypt_are_inverse_aes256cbc() {
        secret_stream_sequential_test_suite(SymKeyKind::Aes256Cbc)
    }

    #[test]
    fn secret_stream_encrypt_decrypt_are_inverse_aes256ctr() {
        secret_stream_sequential_test_suite(SymKeyKind::Aes256Ctr)
    }

    fn secret_stream_random_access_test_case(
        rando: Randomizer,
        key: SymKey,
        inner_sect_sz: usize,
        sect_ct: usize,
    ) {
        let mut stream = SecretStream::new(key)
            .try_compose(SectoredCursor::new(
                vec![0u8; inner_sect_sz * sect_ct],
                inner_sect_sz,
            ))
            .expect("compose failed");
        let sect_sz = stream.sector_sz();
        let indices: Vec<usize> = rando.take(sect_ct).map(|e| e % sect_ct).collect();
        for index in indices.iter().map(|e| *e) {
            let offset = index * sect_sz;
            stream
                .seek(SeekFrom::Start(offset as u64))
                .expect("seek to write failed");
            let sector = vec![index as u8; sect_sz];
            stream.write(&sector).expect("write failed");
        }
        for index in indices.iter().map(|e| *e) {
            let offset = index * sect_sz;
            stream
                .seek(SeekFrom::Start(offset as u64))
                .expect("seek to read failed");
            let expected = vec![index as u8; sect_sz];
            let mut actual = vec![0u8; sect_sz];
            stream.read(&mut actual).expect("read failed");
            assert_eq!(expected, actual);
        }
    }

    fn secret_stream_random_access_test_suite(kind: SymKeyKind) {
        const SEED: [u8; Randomizer::HASH.len()] = [3u8; Randomizer::HASH.len()];
        let key = SymKey::generate(kind).expect("key generation failed");
        secret_stream_random_access_test_case(
            Randomizer::new(SEED),
            key.clone(),
            SECTOR_SZ_DEFAULT,
            20,
        );
        secret_stream_random_access_test_case(
            Randomizer::new(SEED),
            key.clone(),
            SECTOR_SZ_DEFAULT,
            800,
        );
        secret_stream_random_access_test_case(Randomizer::new(SEED), key.clone(), 512, 200);
        secret_stream_random_access_test_case(Randomizer::new(SEED), key.clone(), 512, 20);
        secret_stream_random_access_test_case(Randomizer::new(SEED), key.clone(), 512, 200);
    }

    #[test]
    fn secret_stream_random_access() {
        secret_stream_random_access_test_suite(SymKeyKind::Aes256Cbc);
        secret_stream_random_access_test_suite(SymKeyKind::Aes256Ctr);
    }

    fn make_secret_stream(
        key_kind: SymKeyKind,
        num_sectors: usize,
    ) -> SecretStream<SectoredCursor<Vec<u8>>> {
        let key = SymKey::generate(key_kind).expect("key generation failed");
        let inner = SectoredCursor::new(
            vec![0u8; num_sectors * SECTOR_SZ_DEFAULT],
            SECTOR_SZ_DEFAULT,
        );
        SecretStream::new(key)
            .try_compose(inner)
            .expect("compose failed")
    }

    #[test]
    fn secret_stream_seek_from_start() {
        let mut stream = make_secret_stream(SymKeyKind::Aes256Cbc, 3);
        let sector_sz = stream.sector_sz();
        let expected = vec![2u8; sector_sz];
        // Write one sector of ones, one sector of twos and one sector of threes.
        for k in 1..4 {
            let sector: Vec<u8> = std::iter::repeat(k as u8).take(sector_sz).collect();
            stream.write(&sector).expect("writing to stream failed");
        }

        stream
            .seek(SeekFrom::Start(sector_sz as u64))
            .expect("seek failed");

        // A read from the stream should now return the second sector, which is filled with twos.
        let mut actual = vec![0u8; sector_sz];
        stream
            .read(&mut actual)
            .expect("reading from stream failed");
        assert_eq!(expected, actual);
    }

    #[test]
    fn secret_stream_seek_from_current() {
        let mut stream = make_secret_stream(SymKeyKind::Aes256Cbc, 3);
        let sector_sz = stream.sector_sz();
        let expected = vec![3u8; sector_sz];
        // Write one sector of ones, one sector of twos and one sector of threes.
        for k in 1..4 {
            let sector: Vec<u8> = std::iter::repeat(k as u8).take(sector_sz).collect();
            stream.write(&sector).expect("writing to stream failed");
        }

        stream
            .seek(SeekFrom::Current(-1 * (sector_sz as i64)))
            .expect("seek failed");

        // A read from the stream should now return the last sector, which is filled with threes.
        let mut actual = vec![0u8; sector_sz];
        stream
            .read(&mut actual)
            .expect("reading from stream failed");
        assert_eq!(expected, actual);
    }
}