//Imports
var Trie = require('merkle-patricia-tree');
var rlp = require('rlp');
var levelup = require('levelup');
var leveldown = require('leveldown');
var keccak256 = require('js-sha3').keccak256;

// Chains' path
var dbPathRinkeby = '/home/user/.ethereum/rinkeby/geth/chaindata';
var dbPathMainnet = '/home/user/.ethereum/geth/chaindata';

//Path to chain to use
dbPath = dbPathMainnet;

//Contract address
var contract = '0x5db9b94f23d6722E5C30934B53b352B967fb0F7d';

//Create database
var db = levelup(leveldown(dbPath));

// Can be obtained with web3.eth.getBlock(<blockNumber>).stateRoot
var ROOT = '0x5e27b33d4fa2349b744f7c80c60f5bcdfa9c0754b7f2be61972b220bf2ec4d78';
var trie = new Trie(db, ROOT);

/*
   @dev Will return the contract's storage size in bytes
   @param address Contract address to query storage size.
   @param ROOT Block stateRoot to query storage.
*/
async function getContractStorageSize(address, ROOT){

  //Checking if 0x was provided for address
  if (address.substr(0,2) == '0x'){
    address = address.substr(2);
  }

  //Take hash of address
  var addressHash = keccak256(new Buffer(address, 'hex'));

  // Will count the byte size of the contract's storage
  var contractStorageSize = 0;

  //Promise to return
  return new Promise(function(resolve, reject){

    //Verifying if ROOT exists
    trie.checkRoot(ROOT, function (err, ifExists) {
      if (!ifExists) {
        console.log('State root does not exists');
      }
    });

    //Querying the state trie of the given contract
    trie.get('0x' + addressHash, function (err, val) {
      var decodedVal = rlp.decode(val);

      //Verifying if decoded value is valid
      if (!decodedVal || decodedVal.length < 4) {
        reject('The value for the address must be an array of 4 elements');
      }

      // 3rd element in the array is storage root, 1st - nonce, 2nd - balance, 4th - codeHash
      var storageRoot = decodedVal[2];

      //Set trie root to storageRoot
      trie.root = storageRoot;

      // Verifying if storage root exists
      trie.checkRoot(storageRoot, function (err, ifExists) {
        if (!ifExists) {
          reject('Storage root does not exists');
        };
      });

      // Read all entries from contract storage
      var stream = trie.createReadStream();
      
      stream.on('data', function (data) {
        // Obtain Buffer value for current storage node
        var decodedValNode = rlp.decode(data.value);

        // Update contract Storage size
        contractStorageSize += decodedValNode.byteLength;
      });

      //When stream over
      stream.on('end', async function (val) {
        resolve(contractStorageSize);
      });

    });
  });
};

//Printing the contract storage size for test contract
getContractStorageSize(contract, ROOT).then(
    function(contractSize){
      console.log('Contract address : ', contract);
      console.log('Contract Storage size (Bytes): ', contractSize);
    }
);