from __future__ import annotations

import os
import sys
from typing import Any, Callable, List, TextIO, Union

from phasm.codestyle import phasm_render
from phasm.wasm import (
    WasmTypeFloat32,
    WasmTypeFloat64,
    WasmTypeInt32,
    WasmTypeInt64,
)

from . import memory, runners

DASHES = '-' * 16

class InvalidArgumentException(Exception):
    pass

class SuiteResult:
    def __init__(self) -> None:
        self.returned_value = None

RUNNER_CLASS_MAP = {
    'wasmtime': runners.RunnerWasmtime,
}

TRACE_CODE_PREFIX = """
@imported
def trace_adr(adr: i32) -> i32:
    pass

@imported
def trace_i32(x: i32) -> i32:
    pass

"""

def make_int_trace(prefix: str) -> Callable[[int], int]:
    def trace_helper(adr: int) -> int:
        sys.stderr.write(f'{prefix} {adr}\n')
        sys.stderr.flush()
        return adr

    return trace_helper

class Suite:
    """
    WebAssembly test suite
    """
    def __init__(self, code_py: str) -> None:
        self.code_py = code_py

    def run_code(
            self,
            *args: Any,
            runtime: str = 'wasmtime',
            func_name: str = 'testEntry',
            imports: runners.Imports = None,
            do_format_check: bool = True,
            verbose: bool | None = None,
            with_traces: bool = False,
        ) -> Any:
        """
        Compiles the given python code into wasm and
        then runs it

        Returned is an object with the results set
        """
        if verbose is None:
            verbose = bool(os.environ.get('VERBOSE'))

        code_prefix = ''

        if with_traces:
            assert do_format_check is False

            if imports is None:
                imports = {}

            imports.update({
                'trace_adr': make_int_trace('ADR'),
                'trace_i32': make_int_trace('i32'),
            })

            code_prefix = TRACE_CODE_PREFIX

        class_ = RUNNER_CLASS_MAP[runtime]

        runner = class_(code_prefix + self.code_py)

        if verbose:
            write_header(sys.stderr, 'Phasm')
            runner.dump_phasm_code(sys.stderr)

        runner.parse(verbose=verbose)
        runner.compile_ast()
        runner.optimise_wasm_ast()
        runner.compile_wat()

        if verbose:
            write_header(sys.stderr, 'Assembly')
            runner.dump_wasm_wat(sys.stderr)

        runner.interpreter_setup()
        runner.interpreter_load(imports)

        allocator_generator = memory.Allocator(runner.phasm_ast.build, runner)

        # Check if code formatting works
        if do_format_check:
            assert self.code_py == '\n' + phasm_render(runner.phasm_ast) # \n for formatting in tests

        func = runner.phasm_ast.functions[func_name]
        assert func.type5 is not None  # Type hint
        func_args = runner.phasm_ast.build.type5_is_function(func.type5)
        assert func_args is not None
        func_ret = func_args.pop()
        if len(func_args) != len(args):
            raise RuntimeError(f'Invalid number of args for {func_name}')

        wasm_args: List[Union[float, int]] = []
        if args:
            if verbose:
                write_header(sys.stderr, 'Memory (pre alloc)')
                runner.interpreter_dump_memory(sys.stderr)

            for arg, arg_typ in zip(args, func_args, strict=True):
                arg_typ_info = runner.phasm_ast.build.type_info_map.get(arg_typ.name)

                if arg_typ_info and (arg_typ_info.wasm_type is WasmTypeInt32 or arg_typ_info.wasm_type is WasmTypeInt64):
                    assert isinstance(arg, int)
                    wasm_args.append(arg)
                    continue

                if arg_typ_info and (arg_typ_info.wasm_type is WasmTypeFloat32 or arg_typ_info.wasm_type is WasmTypeFloat64):
                    assert isinstance(arg, float)
                    wasm_args.append(arg)
                    continue

                allocator = allocator_generator(arg_typ)
                adr = allocator(arg)
                wasm_args.append(adr)

        if verbose:
            write_header(sys.stderr, 'Memory (pre run)')
            runner.interpreter_dump_memory(sys.stderr)

        result = SuiteResult()
        result.returned_value = runner.call(func_name, *wasm_args)

        extractor = memory.Extractor(runner.phasm_ast.build, runner)(func_ret)
        result.returned_value = extractor(result.returned_value)

        if verbose:
            write_header(sys.stderr, 'Memory (post run)')
            runner.interpreter_dump_memory(sys.stderr)

        return result

def write_header(textio: TextIO, msg: str) -> None:
    textio.write(f'{DASHES} {msg.ljust(16)}  {DASHES}\n')