// Copyright 2015-2020 Parity Technologies (UK) Ltd.
// This file is part of OpenEthereum.

// OpenEthereum is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// OpenEthereum is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with OpenEthereum.  If not, see <http://www.gnu.org/licenses/>.

//! Snapshot test helpers. These are used to build blockchains and state tries
//! which can be queried before and after a full snapshot/restore cycle.

extern crate trie_standardmap;

use hash::KECCAK_NULL_RLP;
use std::sync::Arc;

use account_db::AccountDBMut;
use blockchain::{BlockChain, BlockChainDB};
use client::{ChainInfo, Client};
use engines::EthEngine;
use snapshot::{
    io::{PackedReader, PackedWriter, SnapshotReader},
    StateRebuilder,
};
use types::basic_account::BasicAccount;

use rand::Rng;
use tempdir::TempDir;

use self::trie_standardmap::{Alphabet, StandardMap, ValueMode};
use ethereum_types::H256;
use ethtrie::{SecTrieDBMut, TrieDB, TrieDBMut};
use hash_db::HashDB;
use journaldb;
use keccak_hasher::KeccakHasher;
use kvdb::DBValue;
use trie::{Trie, TrieMut};

// the proportion of accounts we will alter each tick.
const ACCOUNT_CHURN: f32 = 0.01;

/// This structure will incrementally alter a state given an rng.
pub struct StateProducer {
    state_root: H256,
    storage_seed: H256,
}

impl StateProducer {
    /// Create a new `StateProducer`.
    pub fn new() -> Self {
        StateProducer {
            state_root: KECCAK_NULL_RLP,
            storage_seed: H256::zero(),
        }
    }

    /// Tick the state producer. This alters the state, writing new data into
    /// the database.
    pub fn tick<R: Rng>(&mut self, rng: &mut R, db: &mut dyn HashDB<KeccakHasher, DBValue>) {
        // modify existing accounts.
        let mut accounts_to_modify: Vec<_> = {
            let trie = TrieDB::new(&db, &self.state_root).unwrap();
            let temp = trie
                .iter()
                .unwrap() // binding required due to complicated lifetime stuff
                .filter(|_| rng.gen::<f32>() < ACCOUNT_CHURN)
                .map(Result::unwrap)
                .map(|(k, v)| (H256::from_slice(&k), v.to_owned()))
                .collect();

            temp
        };

        // sweep once to alter storage tries.
        for &mut (ref mut address_hash, ref mut account_data) in &mut accounts_to_modify {
            let mut account: BasicAccount =
                ::rlp::decode(&*account_data).expect("error decoding basic account");
            let acct_db = AccountDBMut::from_hash(db, *address_hash);
            fill_storage(acct_db, &mut account.storage_root, &mut self.storage_seed);
            *account_data = DBValue::from_vec(::rlp::encode(&account));
        }

        // sweep again to alter account trie.
        let mut trie = TrieDBMut::from_existing(db, &mut self.state_root).unwrap();

        for (address_hash, account_data) in accounts_to_modify {
            trie.insert(&address_hash[..], &account_data).unwrap();
        }

        // add between 0 and 5 new accounts each tick.
        let new_accs = rng.gen::<u32>() % 5;

        for _ in 0..new_accs {
            let address_hash = H256(rng.gen());
            let balance: usize = rng.gen();
            let nonce: usize = rng.gen();
            let acc = ::state::Account::new_basic(balance.into(), nonce.into()).rlp();
            trie.insert(&address_hash[..], &acc).unwrap();
        }
    }

    /// Get the current state root.
    pub fn state_root(&self) -> H256 {
        self.state_root
    }
}

/// Fill the storage of an account.
pub fn fill_storage(mut db: AccountDBMut, root: &mut H256, seed: &mut H256) {
    let map = StandardMap {
        alphabet: Alphabet::All,
        min_key: 6,
        journal_key: 6,
        value_mode: ValueMode::Random,
        count: 100,
    };
    {
        let mut trie = if *root == KECCAK_NULL_RLP {
            SecTrieDBMut::new(&mut db, root)
        } else {
            SecTrieDBMut::from_existing(&mut db, root).unwrap()
        };

        for (k, v) in map.make_with(&mut seed.to_fixed_bytes()) {
            trie.insert(&k, &v).unwrap();
        }
    }
}

/// Take a snapshot from the given client into a temporary file.
/// Return a snapshot reader for it.
pub fn snap(client: &Client) -> (Box<dyn SnapshotReader>, TempDir) {
    use types::ids::BlockId;

    let tempdir = TempDir::new("").unwrap();
    let path = tempdir.path().join("file");
    let writer = PackedWriter::new(&path).unwrap();
    let progress = Default::default();

    let hash = client.chain_info().best_block_hash;
    client
        .take_snapshot(writer, BlockId::Hash(hash), &progress)
        .unwrap();

    let reader = PackedReader::new(&path).unwrap().unwrap();

    (Box::new(reader), tempdir)
}

/// Restore a snapshot into a given database. This will read chunks from the given reader
/// write into the given database.
pub fn restore(
    db: Arc<dyn BlockChainDB>,
    engine: &dyn EthEngine,
    reader: &dyn SnapshotReader,
    genesis: &[u8],
) -> Result<(), ::error::Error> {
    use std::sync::atomic::AtomicBool;

    let flag = AtomicBool::new(true);
    let components = engine.snapshot_components().unwrap();
    let manifest = reader.manifest();

    let mut state = StateRebuilder::new(db.key_value().clone(), journaldb::Algorithm::Archive);
    let mut secondary = {
        let chain = BlockChain::new(
            Default::default(),
            genesis,
            db.clone(),
            engine.params().eip1559_transition,
        );
        components.rebuilder(chain, db, manifest).unwrap()
    };

    let mut snappy_buffer = Vec::new();

    trace!(target: "snapshot", "restoring state");
    for state_chunk_hash in manifest.state_hashes.iter() {
        trace!(target: "snapshot", "state chunk hash: {}", state_chunk_hash);
        let chunk = reader.chunk(*state_chunk_hash).unwrap();
        let len = snappy::decompress_into(&chunk, &mut snappy_buffer).unwrap();
        state.feed(&snappy_buffer[..len], &flag)?;
    }

    trace!(target: "snapshot", "restoring secondary");
    for chunk_hash in manifest.block_hashes.iter() {
        let chunk = reader.chunk(*chunk_hash).unwrap();
        let len = snappy::decompress_into(&chunk, &mut snappy_buffer).unwrap();
        secondary.feed(&snappy_buffer[..len], engine, &flag)?;
    }

    trace!(target: "snapshot", "finalizing");
    state.finalize(manifest.block_number, manifest.block_hash)?;
    secondary.finalize(engine)
}