import contextlib
import textwrap

import vyper.ast as vy_ast
import vyper.semantics.analysis as analysis
from vyper.ast.utils import parse_to_ast
from vyper.codegen.core import anchor_opt_level
from vyper.codegen.function_definitions import generate_ir_for_function
from vyper.codegen.ir_node import IRnode
from vyper.evm.opcodes import anchor_evm_version
from vyper.exceptions import InvalidType
from vyper.ir import compile_ir, optimizer
from vyper.semantics.analysis.utils import get_exact_type_from_node

from boa.contracts.vyper.ir_executor import executor_from_ir

# id used internally for method id name
_METHOD_ID_VAR = "_calldata_method_id"


@contextlib.contextmanager
def anchor_compiler_settings(compiler_data):
    settings = compiler_data.settings
    with anchor_opt_level(settings.optimize), anchor_evm_version(settings.evm_version):
        yield


def compile_vyper_function(vyper_function, contract):
    """Compiles a vyper function and appends it to the top of the IR of a
    contract. This is useful for vyper `eval` and internal functions, where
    the runtime bytecode must be changed to add more runtime functionality
    (such as eval, and calling internal functions)
    (performance note: this function is very very slow!)
    """

    compiler_data = contract.compiler_data

    with anchor_compiler_settings(compiler_data):
        global_ctx = contract.global_ctx
        ifaces = compiler_data.interface_codes
        ast = parse_to_ast(vyper_function, ifaces)
        vy_ast.folding.fold(ast)

        # override namespace and add wrapper code at the top
        with contract.override_vyper_namespace():
            analysis.add_module_namespace(ast, ifaces)
            analysis.validate_functions(ast)

        ast = ast.body[0]
        func_t = ast._metadata["type"]

        external_func_info = generate_ir_for_function(ast, global_ctx, False)
        ir = external_func_info.common_ir

        entry_label = func_t._ir_info.external_function_base_entry_label

        ir = ["seq", ["goto", entry_label], ir]

        # use a dummy method id
        ir = ["with", _METHOD_ID_VAR, 0, ir]

        # first mush it with the rest of the IR in the contract to ensure
        # all labels are present, and then optimize all together
        # (use unoptimized IR, ir_executor can't handle optimized selector tables)
        _, contract_runtime = contract.unoptimized_ir
        ir = IRnode.from_list(["seq", ir, contract_runtime])
        ir = optimizer.optimize(ir)

        assembly = compile_ir.compile_to_assembly(ir)
        bytecode, source_map = compile_ir.assembly_to_evm(assembly)
        bytecode += contract.data_section
        typ = func_t.return_type

        # generate the IR executor
        ir_executor = executor_from_ir(ir, compiler_data)

        return ast, ir_executor, bytecode, source_map, typ


def generate_bytecode_for_internal_fn(fn):
    """Wraps internal fns with an external fn and generated bytecode"""

    contract = fn.contract
    fn_name = fn.func_t.name
    fn_args = ", ".join([arg.name for arg in fn.func_t.arguments])

    return_sig = ""
    fn_call = ""
    if fn.func_t.return_type:
        return_sig = f" -> {fn.func_t.return_type}"
        fn_call = "return "
    fn_call += f"self.{fn_name}({fn_args})"

    # same but with defaults, signatures, etc.:
    _fn_sig = []
    for arg in fn.func_t.arguments:
        sig_arg_text = f"{arg.name}: {arg.typ}"

        # check if arg has a default value:
        if arg.name in fn.func_t.default_values:
            default_value = fn.func_t.default_values[arg.name].value
            sig_arg_text += f" = {default_value}"

        _fn_sig.append(sig_arg_text)
    fn_sig = ", ".join(_fn_sig)

    wrapper_code = textwrap.dedent(
        f"""
        @external
        @payable
        def __boa_private_{fn_name}__({fn_sig}){return_sig}:
            {fn_call}
    """
    )
    return compile_vyper_function(wrapper_code, contract)


def generate_bytecode_for_arbitrary_stmt(source_code, contract):
    """Wraps arbitrary stmts with external fn and generates bytecode"""

    ast = parse_to_ast(source_code)
    vy_ast.folding.fold(ast)
    ast = ast.body[0]

    return_sig = ""
    debug_body = source_code

    if isinstance(ast, vy_ast.Expr):
        with contract.override_vyper_namespace():
            try:
                ast_typ = get_exact_type_from_node(ast.value)
                return_sig = f"-> {ast_typ}"
                debug_body = f"return {source_code}"
            except InvalidType:
                pass

    # wrap code in function so that we can easily generate code for it
    wrapper_code = textwrap.dedent(
        f"""
        @external
        @payable
        def __boa_debug__() {return_sig}:
            {debug_body}
    """
    )
    return compile_vyper_function(wrapper_code, contract)
