"""
Python classes for storing the representation of Web Assembly code,
and being able to conver it to Web Assembly Text Format
"""

from typing import Any, Iterable, List, Optional, Tuple, Union

###
### This part is more intermediate code
###

class OurType:
    def to_python(self) -> str:
        raise NotImplementedError(self, 'to_python')

    def to_wasm(self) -> str:
        raise NotImplementedError(self, 'to_wasm')

    def alloc_size(self) -> int:
        raise NotImplementedError(self, 'alloc_size')

class OurTypeNone(OurType):
    pass

class OurTypeBool(OurType):
    pass

class OurTypeInt32(OurType):
    def to_python(self) -> str:
        return 'i32'

    def to_wasm(self) -> str:
        return 'i32'

    def alloc_size(self) -> int:
        return 4

class OurTypeInt64(OurType):
    def to_wasm(self) -> str:
        return 'i64'

    def alloc_size(self) -> int:
        return 8

class OurTypeFloat32(OurType):
    def to_wasm(self) -> str:
        return 'f32'

    def alloc_size(self) -> int:
        return 4

class OurTypeFloat64(OurType):
    def to_wasm(self) -> str:
        return 'f64'

    def alloc_size(self) -> int:
        return 8

class OurTypeVector(OurType):
    """
    A vector is a 128-bit value
    """
    def to_wasm(self) -> str:
        return 'v128'

    def alloc_size(self) -> int:
        return 16

class OurTypeVectorInt32x4(OurTypeVector):
    """
    4 Int32 values in a single vector
    """

class Constant:
    """
    TODO
    """
    def __init__(self, value: Union[None, bool, int, float]) -> None:
        self.value = value

class ClassMember:
    """
    Represents a class member
    """
    def __init__(self, name: str, type_: OurType, offset: int, default: Optional[Constant]) -> None:
        self.name = name
        self.type = type_
        self.offset = offset
        self.default = default

class OurTypeClass(OurType):
    """
    Represents a class
    """
    def __init__(self, name: str, members: List[ClassMember]) -> None:
        self.name = name
        self.members = members

    def to_wasm(self) -> str:
        return 'i32' # WASM uses 32 bit pointers

    def alloc_size(self) -> int:
        return sum(x.type.alloc_size() for x in self.members)

class TupleMember:
    """
    Represents a tuple member
    """
    def __init__(self, type_: OurType, offset: int) -> None:
        self.type = type_
        self.offset = offset

class OurTypeTuple(OurType):
    """
    Represents a tuple
    """
    def __init__(self, members: List[TupleMember]) -> None:
        self.members = members

    def to_python(self) -> str:
        return 'Tuple[' + (
            ', '.join(x.type.to_python() for x in self.members)
        ) + ']'

    def to_wasm(self) -> str:
        return 'i32' # WASM uses 32 bit pointers

    def alloc_size(self) -> int:
        return sum(x.type.alloc_size() for x in self.members)

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, OurTypeTuple):
            raise NotImplementedError

        return (
            len(self.members) == len(other.members)
            and all(
                self.members[x].type == other.members[x].type
                for x in range(len(self.members))
            )
        )

Param = Tuple[str, OurType]

###
## This part is more actual web assembly
###

class Import:
    """
    Represents a Web Assembly import
    """
    def __init__(
            self,
            module: str,
            name: str,
            intname: str,
            params: Iterable[Param],
        ) -> None:
        self.module = module
        self.name = name
        self.intname = intname
        self.params = [*params]
        self.result: OurType = OurTypeNone()

    def generate(self) -> str:
        """
        Generates the text version
        """
        return '(import "{}" "{}" (func ${}{}))'.format(
            self.module,
            self.name,
            self.intname,
            ''.join(' (param {})'.format(x[1]) for x in self.params)
        )

class Statement:
    """
    Represents a Web Assembly statement
    """
    def __init__(self, name: str, *args: str, comment: Optional[str] = None):
        self.name = name
        self.args = args
        self.comment = comment

    def generate(self) -> str:
        """
        Generates the text version
        """
        args = ' '.join(self.args)
        comment = f' ;; {self.comment}' if self.comment else ''

        return f'{self.name} {args}{comment}'

class Function:
    """
    Represents a Web Assembly function
    """
    def __init__(
            self,
            name: str,
            exported: bool,
            params: Iterable[Param],
            locals_: Iterable[Param],
            result: OurType,
            statements: Iterable[Statement],
        ) -> None:
        self.name = name
        self.exported = exported
        self.params = [*params]
        self.locals = [*locals_]
        self.result = result
        self.statements = [*statements]

    def generate(self) -> str:
        """
        Generates the text version
        """
        header = ('(export "{}")' if self.exported else '${}').format(self.name)

        for nam, typ in self.params:
            header += f' (param ${nam} {typ.to_wasm()})'

        if not isinstance(self.result, OurTypeNone):
            header += f' (result {self.result.to_wasm()})'

        for nam, typ in self.locals:
            header += f' (local ${nam} {typ.to_wasm()})'

        return '(func {}\n    {}\n    )'.format(
            header,
            '\n    '.join(x.generate() for x in self.statements),
        )

class ModuleMemory:
    def __init__(self, data: bytes = b'') -> None:
        self.data = data

    def generate(self) -> str:
        return '(memory 1)\n  (data (memory 0) (i32.const 0) "{}")\n'.format(
            ''.join(
                f'\\{x:02x}'
                for x in self.data
            )
        )

class Module:
    """
    Represents a Web Assembly module
    """
    def __init__(self) -> None:
        self.imports: List[Import] = []
        self.functions: List[Function] = []
        self.memory = ModuleMemory(b'\x04') # For ___new_reference___

    def generate(self) -> str:
        """
        Generates the text version
        """
        return '(module\n  {}\n  {}\n  {})\n'.format(
            self.memory.generate(),
            '\n  '.join(x.generate() for x in self.imports),
            '\n  '.join(x.generate() for x in self.functions),
        )