import json
from mythril.analysis.security import get_detection_module_hooks
from mythril.analysis.symbolic import SymExecWrapper
from mythril.analysis.callgraph import generate_graph
from mythril.ethereum.evmcontract import EVMContract
from mythril.solidity.soliditycontract import SolidityContract
from mythril.mythril import Mythril


from mythril.laser.ethereum.state.account import Account
from mythril.laser.ethereum.state.machine_state import MachineState
from mythril.laser.ethereum.state.global_state import GlobalState
from mythril.laser.ethereum import svm
from tests import (
    BaseTestCase,
    TESTDATA_INPUTS_CONTRACTS,
    TESTDATA_OUTPUTS_EXPECTED_LASER_RESULT,
    TESTDATA_OUTPUTS_CURRENT_LASER_RESULT,
)


class LaserEncoder(json.JSONEncoder):
    def default(self, o):
        if getattr(o, "__module__", None) == "z3.z3":
            return str(o)
        return str(o)


def _all_info(laser):
    accounts = {}
    for address, _account in laser.world_state.accounts.items():
        account = _account.as_dict
        account["code"] = account["code"].instruction_list
        account["balance"] = str(account["balance"])
        accounts[address] = account

    nodes = {}
    for uid, node in laser.nodes.items():
        states = []
        for state in node.states:
            if isinstance(state, MachineState):
                states.append(state.as_dict)
            elif isinstance(state, GlobalState):
                environment = state.environment.as_dict
                environment["active_account"] = environment["active_account"].address
                states.append(
                    {
                        "accounts": state.accounts.keys(),
                        "environment": environment,
                        "mstate": state.mstate.as_dict,
                    }
                )

        nodes[uid] = {
            "uid": node.uid,
            "contract_name": node.contract_name,
            "start_addr": node.start_addr,
            "states": states,
            "constraints": node.constraints,
            "function_name": node.function_name,
            "flags": str(node.flags),
        }

    edges = [edge.as_dict for edge in laser.edges]

    return {
        "accounts": accounts,
        "nodes": nodes,
        "edges": edges,
        "total_states": laser.total_states,
        "max_depth": laser.max_depth,
    }


class SVMTestCase(BaseTestCase):
    def setUp(self):
        super(SVMTestCase, self).setUp()
        svm.gbl_next_uid = 0

    def test_laser_result(self):
        for input_file in TESTDATA_INPUTS_CONTRACTS.iterdir():
            if input_file.name in ["weak_random.sol", "environments.sol"]:
                continue
            output_expected = TESTDATA_OUTPUTS_EXPECTED_LASER_RESULT / (
                input_file.name + ".json"
            )
            output_current = TESTDATA_OUTPUTS_CURRENT_LASER_RESULT / (
                input_file.name + ".json"
            )

            disassembly = SolidityContract(
                str(input_file), solc_binary=Mythril._init_solc_binary("0.5.0")
            ).disassembly
            account = Account("0x0000000000000000000000000000000000000000", disassembly)
            accounts = {account.address: account}

            laser = svm.LaserEVM(accounts, max_depth=22, transaction_count=1)
            laser.register_hooks(
                hook_type="post", hook_dict=get_detection_module_hooks()
            )
            laser.sym_exec(account.address)
            laser_info = _all_info(laser)

            output_current.write_text(
                json.dumps(laser_info, cls=LaserEncoder, indent=4)
            )

            if not (output_expected.read_text() == output_expected.read_text()):
                self.found_changed_files(input_file, output_expected, output_current)

        self.assert_and_show_changed_files()

    def runTest(self):

        code = "0x60606040525b603c5b60006010603e565b9050593681016040523660008237602060003683856040603f5a0204f41560545760206000f35bfe5b50565b005b73c3b2ae46792547a96b9f84405e36d0e07edcd05c5b905600a165627a7a7230582062a884f947232ada573f95940cce9c8bfb7e4e14e21df5af4e884941afb55e590029"

        contract = EVMContract(code)
        sym = SymExecWrapper(contract, "0xd0a6E6C543bC68Db5db3A191B171A77407Ff7ccf")

        html = generate_graph(sym)

        self.assertTrue("0 PUSH1 0x60\\n2 PUSH1 0x40\\n4 MSTORE\\n5 JUMPDEST" in html)