// Copyright 2020-2021 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

use crate::vault::{
    crypto_box::{BoxProvider, Decrypt, Encrypt, Key},
    types::{
        transactions::{DataTransaction, RevocationTransaction, SealedBlob, SealedTransaction},
        utils::{BlobId, ChainId, RecordHint, RecordId, VaultId},
    },
};

use serde::{Deserialize, Serialize};

use std::collections::HashMap;

use runtime::GuardedVec;

/// A view over the data inside of a collection of [`Vault`] types.
#[derive(Deserialize, Serialize, Clone, Default)]
pub struct DbView<P: BoxProvider> {
    /// A hashmap of the [`Vault`] types.
    pub vaults: HashMap<VaultId, Vault<P>>,
}

/// A enclave of data that is encrypted under one [`Key`].
#[derive(Deserialize, Serialize, Clone)]
pub struct Vault<P: BoxProvider> {
    key: Key<P>,
    entries: HashMap<ChainId, Record>,
}

/// A bit of data inside of a [`Vault`].
#[derive(Deserialize, Serialize, Clone)]
pub struct Record {
    /// record id.
    id: ChainId,
    /// data transaction metadata.
    data: SealedTransaction,
    /// revocation transaction metadata.
    revoke: Option<SealedTransaction>,
    /// encrypted data in blob format.
    blob: SealedBlob,
}

impl<P: BoxProvider> DbView<P> {
    /// Create a new [`DbView`] to interface with the [`Vault`] types in the database.
    pub fn new() -> DbView<P> {
        let vaults = HashMap::new();

        Self { vaults }
    }

    /// Initialize a new [`Vault`] if it doesn't exist.
    pub fn init_vault(&mut self, key: &Key<P>, vid: VaultId) -> crate::Result<()> {
        self.vaults.entry(vid).or_insert(Vault::init_vault(key)?);

        Ok(())
    }

    /// Write a new record to a [`Vault`]. Will instead update a [`Record`] if it already exists.
    pub fn write(
        &mut self,
        key: &Key<P>,
        vid: VaultId,
        rid: RecordId,
        data: &[u8],
        record_hint: RecordHint,
    ) -> crate::Result<()> {
        if !self.vaults.contains_key(&vid) {
            self.init_vault(key, vid)?;
        }

        self.vaults.entry(vid).and_modify(|vault| {
            vault
                .add_or_update_record(key, rid.0, data, record_hint)
                .expect("unable to write record")
        });

        Ok(())
    }

    /// Lists all of the [`RecordHint`] values and [`RecordId`] values for the given [`Vault`].
    pub fn list_hints_and_ids(&self, key: &Key<P>, vid: VaultId) -> Vec<(RecordId, RecordHint)> {
        let buf: Vec<(RecordId, RecordHint)> = if let Some(vault) = self.vaults.get(&vid) {
            vault.list_hints_and_ids(key)
        } else {
            vec![]
        };

        buf
    }

    /// Check to see if a [`Vault`] contains a [`Record`] through the given [`RecordId`].
    pub fn contains_record(&mut self, key: &Key<P>, vid: VaultId, rid: RecordId) -> bool {
        if let Some(vault) = self.vaults.get(&vid) {
            vault.contains_record(key, rid)
        } else {
            false
        }
    }

    /// Get access the decrypted [`GuardedVec`] of the specified [`Record`].
    pub fn get_guard<F>(&mut self, key: &Key<P>, vid: VaultId, rid: RecordId, f: F) -> crate::Result<()>
    where
        F: FnOnce(GuardedVec<u8>) -> crate::Result<()>,
    {
        if let Some(vault) = self.vaults.get_mut(&vid) {
            let guard = vault.get_guard(key, rid.0)?;

            f(guard)?;
        }

        Ok(())
    }

    /// Access the decrypted [`GuardedVec`] of the specified [`Record`] and place the return value into the second
    /// specified [`Record`]
    #[allow(clippy::too_many_arguments)]
    pub fn exec_proc<F>(
        &mut self,
        key0: &Key<P>,
        vid0: VaultId,
        rid0: RecordId,
        key1: &Key<P>,
        vid1: VaultId,
        rid1: RecordId,
        hint: RecordHint,
        f: F,
    ) -> crate::Result<()>
    where
        F: FnOnce(GuardedVec<u8>) -> crate::Result<Vec<u8>>,
    {
        if let Some(vault) = self.vaults.get_mut(&vid0) {
            let guard = vault.get_guard(key0, rid0.0)?;

            let data = f(guard)?;

            if self.vaults.get(&vid1).is_none() {
                self.init_vault(key1, vid1)?;
            }

            self.write(key1, vid1, rid1, &data, hint)?;
        }

        Ok(())
    }

    /// Add a revocation transaction to the [`Record`]
    pub fn revoke_record(&mut self, key: &Key<P>, vid: VaultId, rid: RecordId) -> crate::Result<()> {
        if let Some(vault) = self.vaults.get_mut(&vid) {
            vault.revoke(key, rid.0)?;
        }

        Ok(())
    }

    /// Garbage collect a [`Vault`]. Deletes any records that contain revocation transactions.
    pub fn garbage_collect_vault(&mut self, key: &Key<P>, vid: VaultId) -> crate::Result<()> {
        if let Some(vault) = self.vaults.get_mut(&vid) {
            if &vault.key == key {
                vault.garbage_collect();
            }
        }

        Ok(())
    }

    /// Clears the entire [`Vault`] from memory.
    pub fn clear(&mut self) -> crate::Result<()> {
        self.vaults.clear();

        Ok(())
    }
}

impl<P: BoxProvider> Vault<P> {
    /// Initialize a new [`Vault`]
    pub fn init_vault(key: &Key<P>) -> crate::Result<Vault<P>> {
        let entries = HashMap::new();

        Ok(Self {
            entries,
            key: key.clone(),
        })
    }

    /// Adds a new [`Record`] to the [`Vault`] if the [`Record`] doesn't already exist. Otherwise, updates the data in
    /// the existing [`Record`] as long as it hasn't been revoked.
    pub fn add_or_update_record(
        &mut self,
        key: &Key<P>,
        id: ChainId,
        data: &[u8],
        record_hint: RecordHint,
    ) -> crate::Result<()> {
        let blob_id = BlobId::random::<P>()?;

        if key == &self.key {
            self.entries
                .entry(id)
                .and_modify(|entry| {
                    entry.update(key, id, data).expect("Unable to update entry");
                })
                .or_insert(Record::new(key, id, blob_id, data, record_hint)?);
        }
        Ok(())
    }

    /// List the [`RecordHint`] values and [`RecordId`] values of the specified [`Vault`].
    pub(crate) fn list_hints_and_ids(&self, key: &Key<P>) -> Vec<(RecordId, RecordHint)> {
        let mut buf: Vec<(RecordId, RecordHint)> = Vec::new();

        if key == &self.key {
            buf = self
                .entries
                .values()
                .into_iter()
                .filter_map(|entry| entry.get_hint_and_id(key))
                .collect();
        }

        buf
    }

    /// Check if the [`Vault`] contains a [`Record`]
    fn contains_record(&self, key: &Key<P>, rid: RecordId) -> bool {
        if key == &self.key {
            self.entries.values().into_iter().any(|entry| entry.check_id(rid))
        } else {
            false
        }
    }

    /// Revokes an [`Record`] by its [`ChainId`].  Does nothing if the [`Record`] doesn't exist.
    pub fn revoke(&mut self, key: &Key<P>, id: ChainId) -> crate::Result<()> {
        if key == &self.key {
            if let Some(entry) = self.entries.get_mut(&id) {
                entry.revoke(key, id)?;
            }
        }

        Ok(())
    }

    /// Gets the decrypted [`GuardedVec`] from the [`Record`]
    pub fn get_guard(&self, key: &Key<P>, id: ChainId) -> crate::Result<GuardedVec<u8>> {
        if key == &self.key {
            if let Some(entry) = self.entries.get(&id) {
                entry.get_blob(key, id)
            } else {
                Ok(GuardedVec::random(0))
            }
        } else {
            Err(crate::Error::DatabaseError("Invalid key.".into()))
        }
    }

    /// Sorts through all of the vault entries and garbage collects any revoked entries.
    pub fn garbage_collect(&mut self) {
        // get the keys of the entries with the revocation transactions.
        let garbage: Vec<ChainId> = self
            .entries
            .iter()
            .filter(|(_, entry)| entry.revoke.is_some())
            .map(|(c, _)| *c)
            .collect();

        // remove the garbage entries from the database.
        garbage.iter().for_each(|c| {
            self.entries.remove(c);
        });
    }
}

impl Record {
    // create a new [`Record`].
    pub fn new<P: BoxProvider>(
        key: &Key<P>,
        id: ChainId,
        blob: BlobId,
        data: &[u8],
        hint: RecordHint,
    ) -> crate::Result<Record> {
        let len = data.len() as u64;
        let dtx = DataTransaction::new(id, len, blob, hint);

        let blob: SealedBlob = data.encrypt(key, blob)?;
        let data = dtx.encrypt(key, id)?;

        Ok(Record {
            id,
            data,
            blob,
            revoke: None,
        })
    }

    /// gets the [`RecordHint`] and [`RecordId`] of the [`Record`].
    fn get_hint_and_id<P: BoxProvider>(&self, key: &Key<P>) -> Option<(RecordId, RecordHint)> {
        if self.revoke.is_none() {
            let tx = self.data.decrypt(key, self.id).expect("Unable to decrypt transaction");

            let tx = tx
                .typed::<DataTransaction>()
                .expect("Failed to convert to data transaction");

            let hint = tx.record_hint;
            let id = RecordId(self.id);

            Some((id, hint))
        } else {
            None
        }
    }

    /// Check to see if a [`RecordId`] pairs with the [`Record`]. Comes back as false if there is a revocation
    /// transaction
    fn check_id(&self, rid: RecordId) -> bool {
        if self.revoke.is_none() {
            rid.0 == self.id
        } else {
            false
        }
    }

    /// Get the blob from this [`Record`].
    fn get_blob<P: BoxProvider>(&self, key: &Key<P>, id: ChainId) -> crate::Result<GuardedVec<u8>> {
        // check if id id and tx id match.
        if self.id == id {
            // check if there is a revocation transaction.
            if self.revoke.is_none() {
                let tx = self.data.decrypt(key, self.id).expect("Unable to decrypt tx");
                let tx = tx
                    .typed::<DataTransaction>()
                    .expect("Failed to cast as data transaction");

                let guarded = GuardedVec::new(tx.len.u64() as usize, |i| {
                    let blob = SealedBlob::from(self.blob.as_ref())
                        .decrypt(key, tx.blob)
                        .expect("Unable to decrypt blob");

                    i.copy_from_slice(blob.as_ref());
                });

                Ok(guarded)
            } else {
                Ok(GuardedVec::zero(0))
            }
        } else {
            Err(crate::Error::DatabaseError(
                "Invalid id for record. Ids must match a valid record.".to_string(),
            ))
        }
    }

    /// Update the data in an existing [`Record`].
    fn update<P: BoxProvider>(&mut self, key: &Key<P>, id: ChainId, new_data: &[u8]) -> crate::Result<()> {
        // check if ids match
        if self.id == id {
            // check if a revocation transaction exists.
            if self.revoke.is_none() {
                // decrypt data transaction.
                let tx = self.data.decrypt(key, self.id)?;
                let tx = tx
                    .typed::<DataTransaction>()
                    .expect("Unable to cast to data transaction");

                // create a new sealed blob with the new_data.
                let blob: SealedBlob = new_data.encrypt(key, tx.blob)?;
                // create a new sealed transaction with the new_data length.
                let dtx = DataTransaction::new(tx.id, new_data.len() as u64, tx.blob, tx.record_hint);
                let data = dtx.encrypt(key, tx.id)?;

                self.blob = blob;
                self.data = data;
            }
        }

        Ok(())
    }

    // add a revocation transaction to the [`Record`].
    fn revoke<P: BoxProvider>(&mut self, key: &Key<P>, id: ChainId) -> crate::Result<()> {
        // check if id and id match.
        if self.id == id {
            // check if revoke transaction already exists.
            if self.revoke.is_none() {
                let revoke = RevocationTransaction::new(self.id);

                self.revoke = Some(revoke.encrypt(key, self.id)?);
            }
        }

        Ok(())
    }
}
