"""
This script generates VMTests that are used to check EVM's Frontier fork correctness.

### TO GENERATE ALL:

## Initialize env:
cd manticore/tests/ && mkdir -p  ethereum_vm/VMTests_concrete && mkdir ethereum_vm/VMTests_symbolic
git clone https://github.com/ethereum/tests --depth=1

## Generate concrete tests:
for i in ./tests/VMTests/*; do python ./auto_generators/make_VMTests.py $i; done

## Generate symbolic tests:
$ for i in ./tests/VMTests/*; do python ./auto_generators/make_VMTests.py $i --symbolic; done

## Remove the eth tests repo
$ rm -rf ./tests  # cleanup/remove the ethereum/tests repo


### To test just one:
python ../auto_generators/make_VMTests.py ./tests/VMTests/VMTests
"""
import json
import os
import sys
from binascii import unhexlify

import hashlib
import pyevmasm as EVMAsm


def gen_test(testcase, filename, symbolic, skip):
    output = """    @unittest.skip('Gas or performance related')\n""" if skip else ""

    # Sanity checks

    # We don't use those!
    testcase.pop("_info")
    assert len(testcase.pop("callcreates", [])) == 0

    output += generate_pre_output(testcase, filename, symbolic)

    if "post" not in testcase:
        output += """
        # If test end in exception check it here
        self.assertTrue(result == 'THROW')"""
    else:
        output += generate_post_output(testcase)

    return output


def generate_pre_output(testcase, filename, symbolic):
    testname = (os.path.split(filename)[1].replace("-", "_")).split(".")[0]
    bytecode = unhexlify(testcase["exec"]["code"][2:])
    disassemble = ""
    try:
        # add silent=True when evmasm supports it
        disassemble = "\n                  ".join(EVMAsm.disassemble(bytecode).split("\n"))
    except Exception as e:
        pass

    with open(filename, "rb") as f:
        sha256sum = hashlib.sha256(f.read()).hexdigest()

    output = f'''
    def test_{testname}(self):
        """
        Testcase taken from https://github.com/ethereum/tests
        File: {os.path.split(filename)[1]}
        sha256sum: {sha256sum}
        Code:     {disassemble}
        """    
    '''

    if symbolic:
        output += """
        def solve(val):
            return self._solve(constraints, val)
"""
    else:
        output += '''
        def solve(val):
            """
            Those tests are **auto-generated** and `solve` is used in symbolic tests.
            So yes, this returns just val; it makes it easier to generate tests like this.
            """
            return to_constant(val)
'''

    env = testcase["env"]

    gaslimit = int(env["currentGasLimit"], 0)
    blocknumber = int(env["currentNumber"], 0)
    timestamp = int(env["currentTimestamp"], 0)
    difficulty = int(env["currentDifficulty"], 0)
    coinbase = int(env["currentCoinbase"], 0)
    output += f"""
        constraints = ConstraintSet()
"""

    def format_bitvec(name, val, bits=256, symbolic_name=None):
        if symbolic_name is None:
            symbolic_name = name

        return f"""
        {name} = constraints.new_bitvec({bits}, name='{symbolic_name}')
        constraints.add({name} == {val})
"""

    def format_var(name, val):
        return f"""
        {name} = {val}"""

    # Spawns/creates bitvectors in symbolic or pure variables in concrete
    formatter = format_bitvec if symbolic else format_var
    for var in ("blocknumber", "timestamp", "difficulty", "coinbase", "gaslimit"):
        output += formatter(var, locals()[var])

    output += f"""
        world = evm.EVMWorld(constraints, blocknumber=blocknumber, timestamp=timestamp, difficulty=difficulty,
                             coinbase=coinbase, gaslimit=gaslimit)
    """

    for address, account in testcase["pre"].items():
        assert account.keys() == {"code", "nonce", "balance", "storage"}

        account_address = int(address, 0)
        account_code = account["code"][2:]
        account_nonce = int(account["nonce"], 0)
        account_balance = int(account["balance"], 0)

        if symbolic:
            postfix = hex(account_address)
            output += f"""
        acc_addr = {hex(account_address)}
        acc_code = unhexlify('{account_code}')
            """
            output += format_bitvec(
                "acc_balance", account_balance, symbolic_name=f"balance_{postfix}"
            )
            output += format_bitvec("acc_nonce", account_nonce, symbolic_name=f"nonce_{postfix}")
        else:
            output += f"""
        acc_addr = {hex(account_address)}
        acc_code = unhexlify('{account_code}')
        acc_balance = {account_balance}
        acc_nonce = {account_nonce}
"""

        output += f"""
        world.create_account(address=acc_addr, balance=acc_balance, code=acc_code, nonce=acc_nonce)
"""
        for key, value in account["storage"].items():
            if symbolic:
                postfix = hex(account_address)
                output += format_bitvec("key", key, symbolic_name=f"storage_key_{postfix}")
                output += format_bitvec("value", value, symbolic_name=f"storage_value_{postfix}")
                output += """
        world.set_storage_data(acc_addr, key, value)
"""
            else:
                output += f"""
        world.set_storage_data(acc_addr, {key}, {value})
"""

    address = int(testcase["exec"]["address"], 0)
    caller = int(testcase["exec"]["caller"], 0)
    code = testcase["exec"]["code"][2:]
    calldata = unhexlify(testcase["exec"]["data"][2:])
    gas = int(testcase["exec"]["gas"], 0)
    price = int(testcase["exec"]["gasPrice"], 0)
    origin = int(testcase["exec"]["origin"], 0)
    value = int(testcase["exec"]["value"], 0)
    assert testcase["exec"].keys() == {
        "address",
        "caller",
        "code",
        "data",
        "gas",
        "gasPrice",
        "origin",
        "value",
    }

    # Need to check if origin is diff from caller. we do not support those tests
    assert origin == caller, "test type not supported"
    assert (
        testcase["pre"]["0x{:040x}".format(address)]["code"][2:] == code
    ), "test type not supported"

    output += f"""
        address = {hex(address)}
        caller = {hex(caller)}"""

    if symbolic:
        output += format_bitvec("price", price)
        output += format_bitvec("value", value)
        output += format_bitvec("gas", gas)
    else:
        output += f"""
        price = {hex(price)}
        value = {value}
        gas = {gas}"""

    if calldata:
        if symbolic: