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

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

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

type SharedCountMap = Arc<Mutex<HashMap<String, usize>>>;
type SharedSet = Arc<Mutex<HashSet<String>>>;

/// Protects checkpoint data by ensuring that no two uncommitted checkpoints can have the same
/// identifier and that committed checkpoints cannot be removed while in use (i.e. have a binding
/// created by either the `commit_checkpoint` method or the `load_checkpoint` method).
///
/// Protection is limited to interactions with a single `GuardWrapper` object directly.
///
/// For example, if multiple guarded [`FileStorage`](
///     https://docs.rs/checkpoint/0.1.5/checkpoint/storage/struct.FileStorage.html) objects
/// are opened with the same path, it is possible for one of them to remove a checkpoint even if
/// the others have the same checkpoint marked as in-use.
///
/// Furthermore, `GuardWrapper` objects offer no protection against outside processes which may be
/// operating on the underlying storage medium.
#[derive(Debug)]
pub struct GuardWrapper<T> {
    inner: T,
    committed_cps_in_use: SharedCountMap,
    uncommitted_cps: SharedSet,
}

impl<T: Storage> GuardWrapper<T> {
    /// Adds a `GuardWrapper` around a `Storage` object.
    pub fn wrap(mut inner: T) -> Result<Self, Error> {
        let committed_cps_in_use = SharedCountMap::default();
        {
            let mut counts = committed_cps_in_use.lock().unwrap();
            let identifiers = inner.checkpoint_identifiers().map_err(|_| {
                Error::initialization(
                    "Could not retrieve existing checkpoint identifiers".to_string(),
                )
            })?;
            for identifier in identifiers {
                counts.insert(identifier, 0);
            }
        }

        Ok(Self {
            inner,
            committed_cps_in_use,
            uncommitted_cps: SharedSet::default(),
        })
    }
}

impl<T: Storage> Storage for GuardWrapper<T> {
    type Committed = Committed<T::Committed>;
    type Uncommitted = Uncommitted<T::Uncommitted>;

    fn create_checkpoint(&mut self, identifier: &str) -> Result<Self::Uncommitted, Error> {
        let committed_cps = self.committed_cps_in_use.lock().unwrap();
        let mut uncommitted_cps = self.uncommitted_cps.lock().unwrap();

        if committed_cps.contains_key(identifier) || uncommitted_cps.contains(identifier) {
            Err(Error::cp_already_exists(format!(
                "Checkpoint {} already exists",
                identifier
            )))?;
        }

        let inner = self.inner.create_checkpoint(identifier)?;
        uncommitted_cps.insert(identifier.to_string());
        Ok(Self::Uncommitted {
            inner: Some(inner),
            identifier: identifier.to_string(),
            uncommitted_cp_set: self.uncommitted_cps.clone(),
        })
    }

    fn commit_checkpoint(
        &mut self,
        mut uncommitted: Self::Uncommitted,
    ) -> Result<Self::Committed, Error> {
        let identifier = uncommitted.identifier()?.to_string();
        let mut committed_cps = self.committed_cps_in_use.lock().unwrap();
        if committed_cps.contains_key(&identifier) {
            Err(Error::cp_already_exists(format!(
                "Checkpoint {} has already been committed",
                &identifier
            )))?;
        }

        let inner = self.inner.commit_checkpoint(uncommitted.take_inner()?)?;
        committed_cps.insert(identifier.clone(), 1);
        Ok(Self::Committed {
            committed_cp_counts: self.committed_cps_in_use.clone(),
            inner,
            identifier,
        })
    }

    fn load_checkpoint(&mut self, identifier: &str) -> Result<Self::Committed, Error> {
        let mut committed_cps = self.committed_cps_in_use.lock().unwrap();
        if !committed_cps.contains_key(identifier) {
            Err(Error::cp_doesnt_exist(format!(
                "Checkpoint {} does not exist",
                identifier
            )))?;
        }

        let inner = self.inner.load_checkpoint(identifier)?;

        *committed_cps.entry(identifier.to_string()).or_insert(0) += 1;
        Ok(Self::Committed {
            committed_cp_counts: self.committed_cps_in_use.clone(),
            inner,
            identifier: identifier.to_string(),
        })
    }

    fn remove_checkpoint(&mut self, identifier: &str) -> Result<(), Error> {
        let mut counts = self.committed_cps_in_use.lock().unwrap();
        if let Some(count) = counts.get(identifier) {
            if *count > 0 {
                Err(Error::cp_in_use(format!(
                    "Checkpoint {} already in use",
                    identifier
                )))?;
            }
        } else {
            Err(Error::cp_doesnt_exist(format!(
                "Checkpoint {} doesn't exist",
                identifier
            )))?;
        }

        self.inner.remove_checkpoint(identifier)?;

        counts.remove(identifier);
        Ok(())
    }

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

/// Wraps the `CommittedCheckpoint` objects from the underlying storage to keep track of which
/// committed checkpoints are in use.
#[derive(Debug)]
pub struct Committed<C: CommittedCheckpoint> {
    committed_cp_counts: SharedCountMap,
    inner: C,
    identifier: String,
}

impl<C: CommittedCheckpoint> CommittedCheckpoint for Committed<C> {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T, Error> {
        self.inner.get(key)
    }

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

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

impl<C: CommittedCheckpoint> Drop for Committed<C> {
    fn drop(&mut self) {
        let identifiers = self.identifier().unwrap().to_string();
        let mut committed_cp_counts = self.committed_cp_counts.lock().unwrap();
        if let Some(count) = committed_cp_counts.get_mut(&identifiers) {
            if *count > 0 {
                *count -= 1;
            }
        }
    }
}

/// Wraps the `UncommittedCheckpoint` objects from the underlying storage to keep track of which
/// uncommitted checkpoints are in use.
#[derive(Debug)]
pub struct Uncommitted<U: UncommittedCheckpoint> {
    inner: Option<U>,
    identifier: String,
    uncommitted_cp_set: SharedSet,
}

impl<U: UncommittedCheckpoint> Uncommitted<U> {
    fn inner(&mut self) -> Result<&mut U, Error> {
        Ok(self.inner.as_mut().ok_or_else(|| {
            Error::other(
                "Checkpoint with type Checkpoint::wrappers::guard_wrapper::Uncommitted has \
                 no inner storage"
                    .to_string(),
            )
        })?)
    }

    fn take_inner(&mut self) -> Result<U, Error> {
        match self.inner.take() {
            Some(i) => Ok(i),
            None => Err(Error::other(
                "Checkpoint with type Checkpoint::wrappers::guard_wrapper::Uncommitted has \
                 no inner storage"
                    .to_string(),
            ))?,
        }
    }
}

impl<U: UncommittedCheckpoint> UncommittedCheckpoint for Uncommitted<U> {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T, Error> {
        self.inner()?.get(key)
    }

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

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

    fn put<T: DeserializeOwned + Serialize>(&mut self, key: &str, value: &T) -> Result<(), Error> {
        self.inner()?.put(key, value)
    }

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

impl<U: UncommittedCheckpoint> Drop for Uncommitted<U> {
    fn drop(&mut self) {
        let identifier = self.identifier().unwrap().to_string();
        let mut uncommitted_cps = self.uncommitted_cp_set.lock().unwrap();
        uncommitted_cps.remove(&identifier);
    }
}

// This module is tested by the tests in file_storage.rs and memory_storage.rs