package io.iohk.ethereum.mpt

import java.io.File

import akka.util.ByteString
import io.iohk.ethereum.crypto._
import io.iohk.ethereum.network.p2p.messages.CommonMessages.{Transaction, TransactionData}
import io.iohk.ethereum.rlp.RLPImplicits._
import io.iohk.ethereum.rlp.{decode => rlpDecode, encode => rlpEncode}
import io.iohk.iodb.LSMStore
import org.spongycastle.util.encoders.Hex

object TestPatriciaTrie {

  implicit val intByteArraySerializable = new ByteArraySerializable[Int] {
    override def toBytes(input: Int): Array[Byte] = rlpEncode(input)

    override def fromBytes(bytes: Array[Byte]): Int = rlpDecode[Int](bytes)
  }

  implicit val txByteArraySerializable = new ByteArraySerializable[Transaction] {
    override def toBytes(input: Transaction): Array[Byte] = rlpEncode(input)

    override def fromBytes(bytes: Array[Byte]): Transaction = rlpDecode[Transaction](bytes)
  }

  case class HashMapDataSource(storage: Map[ByteString, Array[Byte]]) extends DataSource {

    override def get(key: Array[Byte]): Option[Array[Byte]] = storage.get(ByteString(key))

    override def update(rootHash: Array[Byte], toRemove: Seq[Key], toUpdate: Seq[(Key, Value)]): DataSource = {
      val afterUpdate = toUpdate.foldLeft(storage)((storage, toUpdate) => storage + (ByteString(toUpdate._1) -> toUpdate._2))
      HashMapDataSource(afterUpdate)
    }
  }


  def main(args: Array[String]): Unit = {
    val dir = File.createTempFile("iodb", "iodb")
    dir.delete()
    dir.mkdir()

    //open new store
    val dataSource = new IodbDataSource(new LSMStore(dir = dir, keySize = 32))
    val emptyTrie = MerklePatriciaTrie[Int, Transaction](dataSource, (input: Array[Byte]) => sha3(input))

    // Get the same hash as Geth
    val tx = Transaction(
      nonce = 0,
      gasPrice = BigInt(20000000000l),
      gasLimit = BigInt(90000),
      receivingAddress = ByteString(Hex.decode("4dcc858fe9a8048bc778b7170ac29faccd116a4c")),
      value = BigInt(50000000000000000l),
      payload = Right(TransactionData(ByteString.empty)),
      pointSign = 28,
      signatureRandom = ByteString(Hex.decode("5663dbcb92cefdbca6907ed5bef9f5134fbde571726364da691f37462e2c4751")),
      signature = ByteString(Hex.decode("0300f72c91529304101e362913dbaf4ae03caee2fb84ba7d52f5362c103db8a1"))
    )

    val afterInsert = emptyTrie.put(0, tx)
    val hexRoothash = Hex.toHexString(afterInsert.getRootHash)
    println(s"Single root hash: $hexRoothash \n")
    assert(hexRoothash == "98de664ac20c6eb730c0f0eff02e1e1300ecbe907e19dd54b65f3859306190f3")

    // We can create a new trie using the same root and data source
    val sameAsEmptyTrie = MerklePatriciaTrie[Int, Transaction](afterInsert.getRootHash, dataSource, (input: Array[Byte]) => sha3(input))
    println(s"Let's grab the stored tx => ${sameAsEmptyTrie.get(0).get} \n")

    // And you can configure which datasource you would like to use, for example, an inmemory one
    val txs = Seq(
      Transaction(
        nonce = 1,
        gasPrice = BigInt(20000000000l),
        gasLimit = BigInt(90000),
        receivingAddress = ByteString(Hex.decode("4dcc858fe9a8048bc778b7170ac29faccd116a4c")),
        value = BigInt(100000000000000000l),
        payload = Right(TransactionData(ByteString.empty)),
        pointSign = 28,
        signatureRandom = ByteString(Hex.decode("f39a27210ef18b03bea91d6d722db71437d26ad6c85069ade674dba593dd2fe3")),
        signature = ByteString(Hex.decode("644d064f9774c05e51858e33f4bdd440819cc7d278075fb233538047ac9f0e5a"))
      ),
      Transaction(
        nonce = 2,
        gasPrice = BigInt(20000000000l),
        gasLimit = BigInt(90000),
        receivingAddress = ByteString(Hex.decode("4dcc858fe9a8048bc778b7170ac29faccd116a4c")),
        value = BigInt(200000000000000000l),
        payload = Right(TransactionData(ByteString.empty)),
        pointSign = 27,
        signatureRandom = ByteString(Hex.decode("301f8200f4d250cda6ea00701ae339ee73296b9ab1315d1402fcd249aaa0024d")),
        signature = ByteString(Hex.decode("344bd35299031441b984d3a6e3b8884ea2c08483a2aeab69edf9e5c5401edbfa"))
      ), Transaction(
        nonce = 3,
        gasPrice = BigInt(20000000000l),
        gasLimit = BigInt(90000),
        receivingAddress = ByteString(Hex.decode("4dcc858fe9a8048bc778b7170ac29faccd116a4c")),
        value = BigInt(300000000000000000l),
        payload = Right(TransactionData(ByteString.empty)),
        pointSign = 27,
        signatureRandom = ByteString(Hex.decode("9c4df5bf829fecfa6c5932ac30a402898e565f6b45107cb0fd1cd7738f210049")),
        signature = ByteString(Hex.decode("4681d880ebed7b3cfa477023d8ce94314705ed2826a0beac52b121b7d8706b69"))
      )
    )
    val ephemeralDS = HashMapDataSource(Map())
    val trie2 = MerklePatriciaTrie[Int, Transaction](ephemeralDS, (input: Array[Byte]) => sha3(input))

    val toInsert = txs.zipWithIndex

    val trieAfterAllInserts = toInsert.foldLeft(trie2) { (result: MerklePatriciaTrie[Int, Transaction], value) =>
      result.put(value._2, value._1)
    }

    val trieAfterAllInsertsRootHash = Hex.toHexString(trieAfterAllInserts.getRootHash)
    println(s"Trie root hash => ${trieAfterAllInsertsRootHash} \n")
    assert(trieAfterAllInsertsRootHash == "3cf268d513153a048d21ea425412038120f9e60dfd13372c20a5d045dcfbf1e7")
  }

}