// 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 serde::de::DeserializeOwned;
use serde::ser::Serialize;

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

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

type Checkpoint = HashMap<String, Vec<u8>>;
type Checkpoints = Arc<Mutex<HashMap<String, Checkpoint>>>;

/// Provides in-memory checkpoint storage.
#[derive(Debug)]
pub struct MemoryStorage {
    checkpoints: Checkpoints,
}

impl MemoryStorage {
    /// Creates a new `MemoryStorage` wrapped in a
    /// [`GuardWrapper`](
    ///     https://docs.rs/checkpoint/0.1.5/checkpoint/wrappers/struct.GuardWrapper.html).
    pub fn new() -> Result<GuardWrapper<Self>> {
        GuardWrapper::wrap(MemoryStorage {
            checkpoints: Checkpoints::default(),
        })
    }
}

impl Storage for MemoryStorage {
    type Committed = Committed;
    type Uncommitted = Uncommitted;

    fn create_checkpoint(&mut self, identifier: &str) -> Result<Self::Uncommitted> {
        Ok(Self::Uncommitted {
            identifier: identifier.to_string(),
            checkpoint: Checkpoint::default(),
        })
    }

    fn commit_checkpoint(&mut self, uncommitted: Self::Uncommitted) -> Result<Self::Committed> {
        let mut checkpoints = self.checkpoints.lock().unwrap();
        checkpoints.insert(uncommitted.identifier.clone(), uncommitted.checkpoint);

        Ok(Self::Committed {
            identifier: uncommitted.identifier,
            checkpoints: self.checkpoints.clone(),
        })
    }

    fn load_checkpoint(&mut self, identifier: &str) -> Result<Self::Committed> {
        Ok(Self::Committed {
            identifier: identifier.to_string(),
            checkpoints: self.checkpoints.clone(),
        })
    }

    fn remove_checkpoint(&mut self, identifier: &str) -> Result<()> {
        let mut checkpoints = self.checkpoints.lock().unwrap();
        checkpoints.remove(identifier);
        Ok(())
    }

    fn checkpoint_identifiers(&mut self) -> Result<Vec<String>> {
        let committed = self.checkpoints.lock().unwrap();
        let mut identifiers = Vec::new();
        for identifier in committed.keys() {
            identifiers.push(identifier.to_string());
        }

        Ok(identifiers)
    }
}

/// Provides access to committed checkpoint data stored in memory.
#[derive(Debug)]
pub struct Committed {
    identifier: String,
    checkpoints: Checkpoints,
}

impl CommittedCheckpoint for Committed {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T> {
        let checkpoints = self.checkpoints.lock().unwrap();
        let checkpoint = match checkpoints.get(&self.identifier) {
            None => Err(Error::cp_doesnt_exist(format!(
                "Checkpoint {} doesn't exist",
                &self.identifier
            )))?,
            Some(c) => c,
        };

        let data = match checkpoint.get(key) {
            None => Err(Error::cp_key_doesnt_exist(format!(
                "No value stored for checkpoint: {}, key: {}",
                &self.identifier, key
            )))?,
            Some(d) => d,
        };

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

        Ok(value)
    }

    fn keys(&mut self) -> Result<Vec<String>> {
        let checkpoints = self.checkpoints.lock().unwrap();
        let checkpoint = match checkpoints.get(&self.identifier) {
            None => Err(Error::cp_doesnt_exist(format!(
                "Checkpoint {} doesn't exist",
                &self.identifier
            )))?,
            Some(c) => c,
        };

        let mut keys = Vec::new();
        for key in checkpoint.keys() {
            keys.push(key.to_string());
        }

        Ok(keys)
    }

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

/// Provides access to uncommitted checkpoint data stored in memory.
#[derive(Debug)]
pub struct Uncommitted {
    identifier: String,
    checkpoint: Checkpoint,
}

impl UncommittedCheckpoint for Uncommitted {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T> {
        let data = match self.checkpoint.get(key) {
            None => Err(Error::cp_key_doesnt_exist(format!(
                "No value stored for checkpoint: {}, key: {}",
                &self.identifier, key
            )))?,
            Some(d) => d,
        };

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

        Ok(value)
    }

    fn keys(&mut self) -> Result<Vec<String>> {
        let mut keys = Vec::new();
        for key in self.checkpoint.keys() {
            keys.push(key.to_string());
        }

        Ok(keys)
    }

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

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

        self.checkpoint.insert(key.to_string(), data);
        Ok(())
    }

    fn remove(&mut self, key: &str) -> Result<()> {
        if !self.checkpoint.contains_key(key) {
            Err(Error::cp_key_doesnt_exist(format!(
                "No value to remove for checkpoint: {}, key: {}",
                &self.identifier, key
            )))?;
        }

        self.checkpoint.remove(key);
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use std::mem;
    use error::ErrorKind;
    use storage::{CommittedCheckpoint, MemoryStorage, Storage, UncommittedCheckpoint};

    // Verify that checkpoint creation and commit behave as expected.
    #[test]
    fn memory_storage_create_and_commit_checkpoint() {
        let mut storage = MemoryStorage::new().unwrap();

        // Verify that create_checkpoint() can be called twice with the same checkpoint identifier
        // as long as the uncommitted checkpoint returned from the first call is dropped before
        // the second call.
        assert!(storage.create_checkpoint("CP").is_ok());
        assert!(storage.create_checkpoint("CP").is_ok());

        // Verify that create_checkpoint() returns an error when a checkpoint creation is attempted
        // with the same identifier as an uncommitted checkpoint.
        let uncommitted_cp = storage.create_checkpoint("CP").unwrap();
        assert_eq!(
            storage.create_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpAlreadyExists
        );

        // Verify that create_checkpoint() returns an error when a checkpoint creation is attempted
        // with the same identifier as a committed checkpoint that is in use (i.e. has a variable
        // binding).
        let committed_cp = storage.commit_checkpoint(uncommitted_cp).unwrap();
        assert_eq!(
            storage.create_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpAlreadyExists
        );

        // Verify that create_checkpoint() returns an error when a checkpoint creation is attempted
        // with the same identifier as a committed checkpoint that is not in use (i.e. has no
        // variable binding).
        mem::drop(committed_cp);
        assert_eq!(
            storage.create_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpAlreadyExists
        );
    }

    // Verify that checkpoint loading and removal behave as expected.
    #[test]
    fn memory_storage_load_and_remove_checkpoint() {
        let mut storage = MemoryStorage::new().unwrap();

        // Verify that load_checkpoint() returns an error when a checkpoint load is attempted but
        // a committed checkpoint with the given identifier does not exist.
        assert_eq!(
            storage.load_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpDoesntExist
        );

        // Verify that load_checkpoint() returns an error when a checkpoint load is attempted, even
        // if an uncommitted checkpoint with the given identifier exists.
        let uncommitted_cp = storage.create_checkpoint("CP").unwrap();
        assert_eq!(
            storage.load_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpDoesntExist
        );

        // Verify that load_checkpoint() returns a committed checkpoint binding when a committed
        // checkpoint with the given identifier exists and is in use (i.e. has a variable binding).
        let committed_cp = storage.commit_checkpoint(uncommitted_cp).unwrap();
        assert!(storage.load_checkpoint("CP").is_ok());

        // Verify that load_checkpoint() returns a committed checkpoint binding when a committed
        // checkpoint with the given identifier exists and is not in use (i.e. has no
        // variable binding).
        mem::drop(committed_cp);
        let result = storage.load_checkpoint("CP");
        assert!(result.is_ok());

        // Verify that remove_checkpoint() returns an error when an attempt to remove a checkpoint
        // that is in use (i.e. has a variable binding) is made.
        let committed_cp = result.unwrap();
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        // Verify that remove_checkpoint() is successful when an attempt to remove a checkpoint
        // that is not in use (i.e. does not have a variable binding) is made.
        mem::drop(committed_cp);
        assert!(storage.remove_checkpoint("CP").is_ok());

        // Verify that remove_checkpoint() returns an error when an attempt to remove a checkpoint
        // that does not exist is made.
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpDoesntExist
        );

        // Verify that the mechanism which keeps track of the number of times a committed
        // checkpoint has been loaded works as expected.
        let uncommitted_cp = storage.create_checkpoint("CP").unwrap();
        let cp1 = storage.commit_checkpoint(uncommitted_cp).unwrap();
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        let cp2 = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        mem::drop(cp1);
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        let cp3 = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        let cp4 = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        let cp5 = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        // Start dropping loaded checkpoints in an arbitrary order.
        mem::drop(cp4);
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        mem::drop(cp2);
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        mem::drop(cp5);
        assert_eq!(
            storage.remove_checkpoint("CP").unwrap_err().kind(),
            ErrorKind::CpInUse
        );

        mem::drop(cp3);
        assert!(storage.remove_checkpoint("CP").is_ok());
    }

    // Verify that checkpoint identifiers are returned properly.
    #[test]
    fn memory_storage_checkpoint_identifiers() {
        let mut storage = MemoryStorage::new().unwrap();

        // Verify that checkpoint_identifiers() returns an empty list when there are no
        // committed checkpoints.
        assert_eq!(
            storage.checkpoint_identifiers().unwrap(),
            Vec::<String>::new()
        );

        // Verify that checkpoint_identifiers() returns an empty list even when there are
        // uncommitted checkpoints.
        let uncommitted_cp_1 = storage.create_checkpoint("CP1").unwrap();
        assert_eq!(
            storage.checkpoint_identifiers().unwrap(),
            Vec::<String>::new()
        );

        // Verify that checkpoint_identifiers() returns a list of identifiers when committed
        // checkpoints exist and are loaded (i.e. when they have variable bindings).
        let committed_cp_1 = storage.commit_checkpoint(uncommitted_cp_1).unwrap();
        assert_eq!(
            storage.checkpoint_identifiers().unwrap(),
            vec!["CP1".to_string()]
        );

        let uncommitted_cp_2 = storage.create_checkpoint("CP2").unwrap();
        let committed_cp_2 = storage.commit_checkpoint(uncommitted_cp_2).unwrap();
        let mut identifiers = storage.checkpoint_identifiers().unwrap();
        identifiers.sort();
        assert_eq!(identifiers, vec!["CP1".to_string(), "CP2".to_string()]);

        let uncommitted_cp_3 = storage.create_checkpoint("CP3").unwrap();
        let committed_cp_3 = storage.commit_checkpoint(uncommitted_cp_3).unwrap();
        let mut identifiers = storage.checkpoint_identifiers().unwrap();
        identifiers.sort();
        assert_eq!(
            identifiers,
            vec!["CP1".to_string(), "CP2".to_string(), "CP3".to_string()]
        );

        // Verify that checkpoint_identifiers() returns a list of identifiers when committed
        // checkpoints exist and are not loaded (i.e. when they do not have variable bindings).
        mem::drop(committed_cp_1);
        mem::drop(committed_cp_2);
        mem::drop(committed_cp_3);

        let mut identifiers = storage.checkpoint_identifiers().unwrap();
        identifiers.sort();
        assert_eq!(
            identifiers,
            vec!["CP1".to_string(), "CP2".to_string(), "CP3".to_string()]
        );

        // Verify that checkpoint identifiers are removed from the list of identifiers returned by
        // checkpoint_identifiers() when the associated checkpoint is removed.
        storage.remove_checkpoint("CP2").unwrap();
        let mut identifiers = storage.checkpoint_identifiers().unwrap();
        identifiers.sort();
        assert_eq!(identifiers, vec!["CP1".to_string(), "CP3".to_string()]);

        storage.remove_checkpoint("CP1").unwrap();
        let mut identifiers = storage.checkpoint_identifiers().unwrap();
        identifiers.sort();
        assert_eq!(identifiers, vec!["CP3".to_string()]);

        storage.remove_checkpoint("CP3").unwrap();
        let mut identifiers = storage.checkpoint_identifiers().unwrap();
        identifiers.sort();
        assert_eq!(identifiers, Vec::<String>::new());
    }

    // Verify that uncommitted and committed checkpoints behave as expected.
    #[test]
    fn uncommitted_and_committed() {
        let mut storage = MemoryStorage::new().unwrap();

        // Verify that the methods of an empty uncommitted checkpoint behave as expected.
        let mut uncommitted_cp_1 = storage.create_checkpoint("CP1").unwrap();
        assert_eq!(uncommitted_cp_1.identifier().unwrap(), "CP1");
        assert_eq!(uncommitted_cp_1.keys().unwrap(), Vec::<String>::new());
        assert_eq!(
            uncommitted_cp_1.get::<u8>("Key1").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );
        assert_eq!(
            uncommitted_cp_1.remove("Key1").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );

        // Add data to the uncommitted checkpoint.
        let data_1 = 0u8;
        let data_2 = String::from("Test String.");
        let data_3 = Some(true);
        assert!(uncommitted_cp_1.put("Key1", &data_1).is_ok());
        assert!(uncommitted_cp_1.put("Key2", &data_2).is_ok());
        assert!(uncommitted_cp_1.put("Key3", &data_3).is_ok());

        // Verify that adding and removing values to existing checkpoints works as expected and
        // that creating a new uncommitted checkpoint does not affect the data stored in another
        // uncommitted checkpoint.
        let mut uncommitted_cp_2 = storage.create_checkpoint("CP2").unwrap();
        assert_eq!(uncommitted_cp_2.identifier().unwrap(), "CP2");
        assert_eq!(uncommitted_cp_2.keys().unwrap(), Vec::<String>::new());
        assert_eq!(
            uncommitted_cp_2.get::<u8>("Key1").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );
        assert_eq!(
            uncommitted_cp_2.remove("Key1").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );

        assert_eq!(uncommitted_cp_1.identifier().unwrap(), "CP1");
        let mut keys = uncommitted_cp_1.keys().unwrap();
        keys.sort();
        assert_eq!(keys, vec!["Key1", "Key2", "Key3"]);
        assert_eq!(uncommitted_cp_1.get::<u8>("Key1").unwrap(), data_1);
        assert_eq!(uncommitted_cp_1.get::<String>("Key2").unwrap(), data_2);
        assert_eq!(
            uncommitted_cp_1.get::<Option<bool>>("Key3").unwrap(),
            data_3
        );
        assert!(uncommitted_cp_1.remove("Key2").is_ok());
        assert_eq!(
            uncommitted_cp_1.remove("Key2").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );
        let mut keys = uncommitted_cp_1.keys().unwrap();
        keys.sort();
        assert_eq!(keys, vec!["Key1", "Key3"]);
        assert!(uncommitted_cp_1.remove("Key1").is_ok());
        assert_eq!(
            uncommitted_cp_1.remove("Key1").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );
        let mut keys = uncommitted_cp_1.keys().unwrap();
        keys.sort();
        assert_eq!(keys, vec!["Key3"]);
        assert!(uncommitted_cp_1.remove("Key3").is_ok());
        assert_eq!(
            uncommitted_cp_1.remove("Key3").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );
        assert_eq!(uncommitted_cp_1.keys().unwrap(), Vec::<String>::new());

        // Verify that the value for a given key can be added even if the key was previously
        // removed.
        assert!(uncommitted_cp_1.put("Key1", &data_1).is_ok());
        assert!(uncommitted_cp_1.put("Key2", &data_2).is_ok());
        assert!(uncommitted_cp_1.put("Key3", &data_3).is_ok());
        assert_eq!(uncommitted_cp_1.identifier().unwrap(), "CP1");
        let mut keys = uncommitted_cp_1.keys().unwrap();
        keys.sort();
        assert_eq!(keys, vec!["Key1", "Key2", "Key3"]);
        assert_eq!(uncommitted_cp_1.get::<u8>("Key1").unwrap(), data_1);
        assert_eq!(uncommitted_cp_1.get::<String>("Key2").unwrap(), data_2);
        assert_eq!(
            uncommitted_cp_1.get::<Option<bool>>("Key3").unwrap(),
            data_3
        );

        // Verify that the value for a given key can be changed.
        let data_1 = String::from("Test String.");
        let data_2 = Some(true);
        let data_3 = 0u8;
        assert!(uncommitted_cp_1.put("Key1", &data_1).is_ok());
        assert!(uncommitted_cp_1.put("Key2", &data_2).is_ok());
        assert!(uncommitted_cp_1.put("Key3", &data_3).is_ok());
        assert!(uncommitted_cp_1.put("Key1", &data_1).is_ok());
        assert!(uncommitted_cp_1.put("Key2", &data_2).is_ok());
        assert!(uncommitted_cp_1.put("Key3", &data_3).is_ok());
        assert_eq!(uncommitted_cp_1.identifier().unwrap(), "CP1");
        let mut keys = uncommitted_cp_1.keys().unwrap();
        keys.sort();
        assert_eq!(keys, vec!["Key1", "Key2", "Key3"]);
        assert_eq!(uncommitted_cp_1.get::<String>("Key1").unwrap(), data_1);
        assert_eq!(
            uncommitted_cp_1.get::<Option<bool>>("Key2").unwrap(),
            data_2
        );
        assert_eq!(uncommitted_cp_1.get::<u8>("Key3").unwrap(), data_3);

        // Verify that uncommitted checkpoints contain the same values after being committed.
        let mut committed_cp_1 = storage.commit_checkpoint(uncommitted_cp_1).unwrap();
        let mut committed_cp_2 = storage.commit_checkpoint(uncommitted_cp_2).unwrap();

        assert_eq!(committed_cp_1.identifier().unwrap(), "CP1");
        let mut keys = committed_cp_1.keys().unwrap();
        keys.sort();
        assert_eq!(keys, vec!["Key1", "Key2", "Key3"]);
        assert_eq!(committed_cp_1.get::<String>("Key1").unwrap(), data_1);
        assert_eq!(committed_cp_1.get::<Option<bool>>("Key2").unwrap(), data_2);
        assert_eq!(committed_cp_1.get::<u8>("Key3").unwrap(), data_3);
        assert_eq!(
            committed_cp_1.get::<u8>("Key4").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );

        assert_eq!(committed_cp_2.identifier().unwrap(), "CP2");
        assert_eq!(committed_cp_2.keys().unwrap(), Vec::<String>::new());
        assert_eq!(
            committed_cp_2.get::<u8>("Key1").unwrap_err().kind(),
            ErrorKind::CpKeyDoesntExist
        );
    }
}