from typing import Any, Protocol

import struct

from phasm.build.base import BuildBase
from phasm.build.typerouter import BuildTypeRouter
from phasm.type5.typeexpr import AtomicType, TypeExpr
from phasm.wasm import (
    WasmTypeFloat32,
    WasmTypeFloat64,
    WasmTypeInt32,
    WasmTypeInt64,
    WasmTypeNone,
)

class MemoryAccess(Protocol):
    def interpreter_read_memory(self, offset: int, length: int) -> bytes:
        pass

class LoaderFunc(Protocol):
    alloc_size: int

    def __call__(self, wasm_value: Any) -> Any:
        """
        Takes a WASM value and returns a Python value
        Based on the phasm type
        """

class Loader(BuildTypeRouter[LoaderFunc]):
    __slots__ = ('access', )

    access: MemoryAccess

    def __init__(self, build: BuildBase[Any], access: MemoryAccess) -> None:
        super().__init__(build)
        self.access = access

    def when_atomic(self, typ: AtomicType) -> LoaderFunc:
        type_info = self.build.type_info_map[typ.name]

        if type_info.wasm_type is WasmTypeNone:
            raise NotImplementedError()

        if type_info.wasm_type is WasmTypeInt32 or type_info.wasm_type is WasmTypeInt64:
            if type_info.signed is None:
                return BoolLoader()

            return IntLoader(type_info.signed, type_info.alloc_size)

        if type_info.wasm_type is WasmTypeFloat32 or type_info.wasm_type is WasmTypeFloat64:
            return FloatLoader(type_info.alloc_size)

        raise NotImplementedError(typ)

    def when_dynamic_array(self, da_arg: TypeExpr) -> LoaderFunc:
        if da_arg.name == 'u8':
            return BytesLoader(self.access)

        return DynamicArrayLoader(self.access, self(da_arg))

class BoolLoader:
    __slots__ = ('alloc_size', )

    def __init__(self) -> None:
        self.alloc_size = 1

    def __call__(self, wasm_value: Any) -> bool:
        assert isinstance(wasm_value, int), wasm_value

        return wasm_value != 0

class IntLoader:
    __slots__ = ('alloc_size', 'signed', )

    def __init__(self, signed: bool, alloc_size: int) -> None:
        self.signed = signed
        self.alloc_size = alloc_size

    def __call__(self, wasm_value: Any) -> int:
        assert isinstance(wasm_value, int), wasm_value

        # Work around the fact that phasm has unsigned integers but wasm does not
        data = wasm_value.to_bytes(8, 'big', signed=True)
        data = data[-self.alloc_size:]
        return int.from_bytes(data, 'big', signed=self.signed)

class FloatLoader:
    __slots__ = ('alloc_size', )

    def __init__(self, alloc_size: int) -> None:
        self.alloc_size = alloc_size

    def __call__(self, wasm_value: Any) -> float:
        assert isinstance(wasm_value, float), wasm_value
        return wasm_value

class DynamicArrayLoader:
    __slots__ = ('access', 'alloc_size', 'sub_loader', )

    access: MemoryAccess
    alloc_size: int
    sub_loader: LoaderFunc

    def __init__(self, access: MemoryAccess, sub_loader: LoaderFunc) -> None:
        self.access = access
        self.sub_loader = sub_loader

    def __call__(self, wasm_value: Any) -> Any:
        assert isinstance(wasm_value, int), wasm_value
        adr = wasm_value
        del wasm_value

        # wasm_value must be a pointer
        # The first value at said pointer is the length of the array

        read_bytes = self.access.interpreter_read_memory(adr, 4)
        array_len,  = struct.unpack('<I', read_bytes)
        adr += 4

        raise NotImplementedError

class BytesLoader:
    __slots__ = ('access', 'alloc_size', )

    access: MemoryAccess
    alloc_size: int

    def __init__(self, access: MemoryAccess) -> None:
        self.access = access

    def __call__(self, wasm_value: Any) -> bytes:
        assert isinstance(wasm_value, int), wasm_value
        adr = wasm_value
        del wasm_value

        # wasm_value must be a pointer
        # The first value at said pointer is the length of the array

        read_bytes = self.access.interpreter_read_memory(adr, 4)
        array_len,  = struct.unpack('<I', read_bytes)
        adr += 4

        return self.access.interpreter_read_memory(adr, array_len)