// Copyright 2018 Zach Miller
//
// Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
// or the MIT License (https://opensource.org/licenses/MIT), at your option.
//
// This file may not be copied, modified, or distributed except according to those terms.

use bincode;
use sha2::{Digest, Sha256};

use serde::de::DeserializeOwned;
use serde::ser::Serialize;

use error::{Error, Result};
use storage::{CommittedCheckpoint, Storage, UncommittedCheckpoint};

/// Adds checksum verification to checkpoint data.
///
/// Currently the Sha256 checksum algorithm is used.
#[derive(Debug)]
pub struct ChecksumWrapper<S: Storage> {
    inner: S,
}

impl<S: Storage> ChecksumWrapper<S> {
    /// Adds a `ChecksumWrapper` around a `Storage` object.
    pub fn wrap(inner: S) -> Self {
        Self { inner }
    }
}

impl<S: Storage> Storage for ChecksumWrapper<S> {
    type Committed = Committed<S::Committed>;
    type Uncommitted = Uncommitted<S::Uncommitted>;

    fn create_checkpoint(&mut self, identifier: &str) -> Result<Self::Uncommitted> {
        let inner = self.inner.create_checkpoint(identifier)?;
        Ok(Self::Uncommitted { inner })
    }

    fn commit_checkpoint(&mut self, uncommitted: Self::Uncommitted) -> Result<Self::Committed> {
        let inner = self.inner.commit_checkpoint(uncommitted.inner)?;
        Ok(Self::Committed { inner })
    }

    fn load_checkpoint(&mut self, identifier: &str) -> Result<Self::Committed> {
        let inner = self.inner.load_checkpoint(identifier)?;
        Ok(Self::Committed { inner })
    }

    fn remove_checkpoint(&mut self, identifier: &str) -> Result<()> {
        self.inner.remove_checkpoint(identifier)
    }

    fn checkpoint_identifiers(&mut self) -> Result<Vec<String>> {
        self.inner.checkpoint_identifiers()
    }
}

/// Wraps the `CommittedCheckpoint` objects from the underlying storage to verify checkpoint data
/// when it is retrieved.
#[derive(Debug)]
pub struct Committed<C: CommittedCheckpoint> {
    inner: C,
}

impl<C: CommittedCheckpoint> CommittedCheckpoint for Committed<C> {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T> {
        let helper: Helper = self.inner.get(key)?;
        let identifier = self.identifier()?;
        let mut hasher = Sha256::default();
        hasher.input(identifier.as_bytes());
        hasher.input(&helper.bytes);
        if hasher.result()[..] != helper.checksum[..] {
            Err(Error::integrity(format!(
                "Integrity check failed for checkpoint: {}, key: {}",
                identifier, key
            )))?;
        }

        let value = bincode::deserialize(&helper.bytes).map_err(|_| {
            Error::deserialize(format!(
                "Failed to deserialize value from checkpoint: {}, key: {}",
                identifier, key
            ))
        })?;

        Ok(value)
    }

    fn keys(&mut self) -> Result<Vec<String>> {
        self.inner.keys()
    }

    fn identifier(&mut self) -> Result<&str> {
        self.inner.identifier()
    }
}

/// Wraps the `UncommittedCheckpoint` objects from the underlying storage to add a checksum when
/// checkpoint data is stored and to verify checkpoint data when it is retrieved.
#[derive(Debug)]
pub struct Uncommitted<U: UncommittedCheckpoint> {
    inner: U,
}

impl<U: UncommittedCheckpoint> UncommittedCheckpoint for Uncommitted<U> {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T> {
        let helper: Helper = self.inner.get(key)?;
        let identifier = self.identifier()?;
        let mut hasher = Sha256::default();
        hasher.input(identifier.as_bytes());
        hasher.input(&helper.bytes);
        if hasher.result()[..] != helper.checksum[..] {
            Err(Error::integrity(format!(
                "Integrity check failed for checkpoint: {}, key: {}",
                identifier, key
            )))?;
        }

        let value = bincode::deserialize(&helper.bytes).map_err(|_| {
            Error::deserialize(format!(
                "Failed to deserialize value from checkpoint: {}, key: {}",
                identifier, key
            ))
        })?;

        Ok(value)
    }

    fn keys(&mut self) -> Result<Vec<String>> {
        self.inner.keys()
    }

    fn identifier(&mut self) -> Result<&str> {
        self.inner.identifier()
    }

    fn put<T: DeserializeOwned + Serialize>(&mut self, key: &str, value: &T) -> Result<()> {
        let identifier = self.identifier()?.to_string();
        let bytes = bincode::serialize(value).map_err(|_| {
            Error::serialize(format!(
                "Failed to serialize value into checkpoint: {}, key: {}",
                identifier, key
            ))
        })?;

        let mut hasher = Sha256::default();
        hasher.input(identifier.as_bytes());
        hasher.input(&bytes);
        let helper = Helper {
            bytes,
            checksum: hasher.result().to_vec(),
        };

        self.inner.put(key, &helper)
    }

    fn remove(&mut self, key: &str) -> Result<()> {
        self.inner.remove(key)
    }
}

// Used to add a checksum to the values being inserted into the checkpoint.
#[derive(Debug, Deserialize, Serialize)]
struct Helper {
    bytes: Vec<u8>,
    checksum: Vec<u8>,
}

#[cfg(all(test, feature = "filestorage-tests"))]
mod test {
    use std::fs::File;
    use std::io::{Read, Write};

    use error::ErrorKind;
    use storage::{CommittedCheckpoint, FileStorage, Storage, UncommittedCheckpoint};
    use test_utils::create_test_directory_with_checkpoint_folders;
    use wrappers::ChecksumWrapper;

    // Verify that checkpoint data can be read properly when data corruption does not occur.
    #[test]
    fn no_corruption() {
        let dir = create_test_directory_with_checkpoint_folders();
        let path = dir.path();

        let mut storage = ChecksumWrapper::wrap(FileStorage::open(&path).unwrap());
        let mut uncommitted = storage.create_checkpoint("CP").unwrap();
        uncommitted.put("Key", &"Test String".to_string()).unwrap();
        storage.commit_checkpoint(uncommitted).unwrap();

        let mut file_path = path.join("committed");
        file_path.push("CP");
        file_path.push("Key");
        let mut loaded = storage.load_checkpoint("CP").unwrap();
        assert_eq!(loaded.get::<String>("Key").unwrap(), "Test String");
    }

    // Verify that corruption of checkpoint data is detected properly.
    #[test]
    fn corruption_in_stored_data() {
        let dir = create_test_directory_with_checkpoint_folders();
        let path = dir.path();

        let mut storage = ChecksumWrapper::wrap(FileStorage::open(&path).unwrap());
        let mut uncommitted = storage.create_checkpoint("CP").unwrap();
        uncommitted.put("Key", &"Test String".to_string()).unwrap();
        storage.commit_checkpoint(uncommitted).unwrap();

        let mut file_path = path.join("committed");
        file_path.push("CP");
        file_path.push("Key");

        let mut bytes = Vec::new();
        let mut file = File::open(&file_path).unwrap();
        file.read_to_end(&mut bytes).unwrap();
        bytes[16] = b'?';

        let mut file = File::create(&file_path).unwrap();
        file.write(&bytes).unwrap();

        let mut loaded = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            loaded.get::<String>("Key").unwrap_err().kind(),
            ErrorKind::Integrity
        );
    }

    // Verify that corruption of the checksum is detected properly.
    #[test]
    fn corruption_in_checksum() {
        let dir = create_test_directory_with_checkpoint_folders();
        let path = dir.path();

        let mut storage = ChecksumWrapper::wrap(FileStorage::open(&path).unwrap());
        let mut uncommitted = storage.create_checkpoint("CP").unwrap();
        uncommitted.put("Key", &"Test String".to_string()).unwrap();
        storage.commit_checkpoint(uncommitted).unwrap();

        let mut file_path = path.join("committed");
        file_path.push("CP");
        file_path.push("Key");

        let mut bytes = Vec::new();
        let mut file = File::open(&file_path).unwrap();
        file.read_to_end(&mut bytes).unwrap();
        bytes[35] = b'?';

        let mut file = File::create(&file_path).unwrap();
        file.write(&bytes).unwrap();

        let mut loaded = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            loaded.get::<String>("Key").unwrap_err().kind(),
            ErrorKind::Integrity
        );
    }
}