from typing import Any

import struct
import sys

from phasm.codestyle import phasm_render
from phasm.type3 import types as type3types
from phasm.runtime import calculate_alloc_size

from . import runners

DASHES = '-' * 16

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

RUNNER_CLASS_MAP = {
    'pywasm': runners.RunnerPywasm,
    'pywasm3': runners.RunnerPywasm3,
    'wasmtime': runners.RunnerWasmtime,
    'wasmer': runners.RunnerWasmer,
}

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

    def run_code(self, *args, runtime='pywasm3', func_name='testEntry', imports=None):
        """
        Compiles the given python code into wasm and
        then runs it

        Returned is an object with the results set
        """
        class_ = RUNNER_CLASS_MAP[runtime]

        runner = class_(self.code_py)

        runner.parse()
        runner.compile_ast()
        runner.compile_wat()
        runner.compile_wasm()
        runner.interpreter_setup()
        runner.interpreter_load(imports)

        write_header(sys.stderr, 'Phasm')
        runner.dump_phasm_code(sys.stderr)
        write_header(sys.stderr, 'Assembly')
        runner.dump_wasm_wat(sys.stderr)

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

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

            for arg in args:
                if isinstance(arg, (int, float, )):
                    wasm_args.append(arg)
                    continue

                if isinstance(arg, bytes):
                    adr = runner.call('stdlib.types.__alloc_bytes__', len(arg))
                    sys.stderr.write(f'Allocation 0x{adr:08x} {repr(arg)}\n')

                    runner.interpreter_write_memory(adr + 4, arg)
                    wasm_args.append(adr)
                    continue

                raise NotImplementedError(arg)

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

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

        result.returned_value = _load_memory_stored_returned_value(
            runner,
            func_name,
            result.returned_value,
        )

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

        return result

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

def _load_memory_stored_returned_value(
        runner: runners.RunnerBase,
        func_name: str,
        wasm_value: Any,
    ) -> Any:
    ret_type3 = runner.phasm_ast.functions[func_name].returns_type3

    if ret_type3 in (type3types.i32, type3types.i64):
        assert isinstance(wasm_value, int), wasm_value
        return wasm_value

    if ret_type3 in (type3types.u32, type3types.u64):
        assert isinstance(wasm_value, int), wasm_value
        assert 0 <= wasm_value, 'TODO: Extract negative values'
        return wasm_value

    if ret_type3 in (type3types.f32, type3types.f64, ):
        assert isinstance(wasm_value, float), wasm_value
        return wasm_value

    if ret_type3 is type3types.bytes:
        assert isinstance(wasm_value, int), wasm_value
        adr = wasm_value

        sys.stderr.write(f'Reading 0x{adr:08x}\n')
        read_bytes = runner.interpreter_read_memory(adr, 4)
        bytes_len,  = struct.unpack('<I', read_bytes)

        adr += 4
        sys.stderr.write(f'Reading 0x{adr:08x}\n')
        return runner.interpreter_read_memory(adr, bytes_len)

    if isinstance(ret_type3, type3types.AppliedType3):
        if ret_type3.base is type3types.static_array:
            assert isinstance(wasm_value, int), wasm_value
            adr = wasm_value

            assert ret_type3.args[0] is type3types.u64, 'Not Implemented yet'
            assert isinstance(ret_type3.args[1], type3types.IntType3)

            sa_len = ret_type3.args[1].value

            alloc_size = calculate_alloc_size(ret_type3)
            read_bytes = runner.interpreter_read_memory(adr, alloc_size)
            print('read_bytes', read_bytes)

            return struct.unpack(f'<{sa_len}q', read_bytes)

    raise NotImplementedError(ret_type3, wasm_value)