"""
Contains the syntax tree for ourlang
"""
from typing import Dict, List, Optional, Union

import enum

from typing_extensions import Final

WEBASSEMBLY_BUILTIN_FLOAT_OPS: Final = ('abs', 'sqrt', 'ceil', 'floor', 'trunc', 'nearest', )
WEBASSEMBLY_BUILTIN_BYTES_OPS: Final = ('len', )

from .type3 import types as type3types
from .type3.types import Type3, Type3OrPlaceholder, PlaceholderForType, StructType3

class Expression:
    """
    An expression within a statement
    """
    __slots__ = ('type3', )

    type3: Type3OrPlaceholder

    def __init__(self) -> None:
        self.type3 = PlaceholderForType([self])

class Constant(Expression):
    """
    An constant value expression within a statement

    # FIXME: Rename to literal
    """
    __slots__ = ()

class ConstantPrimitive(Constant):
    """
    An primitive constant value expression within a statement
    """
    __slots__ = ('value', )

    value: Union[int, float]

    def __init__(self, value: Union[int, float]) -> None:
        super().__init__()
        self.value = value

    def __repr__(self) -> str:
        return f'ConstantPrimitive({repr(self.value)})'

class ConstantTuple(Constant):
    """
    A Tuple constant value expression within a statement
    """
    __slots__ = ('value', )

    value: List[ConstantPrimitive]

    def __init__(self, value: List[ConstantPrimitive]) -> None: # FIXME: Tuple of tuples?
        super().__init__()
        self.value = value

class VariableReference(Expression):
    """
    An variable reference expression within a statement
    """
    __slots__ = ('variable', )

    variable: Union['ModuleConstantDef', 'FunctionParam'] # also possibly local

    def __init__(self, variable: Union['ModuleConstantDef', 'FunctionParam']) -> None:
        super().__init__()
        self.variable = variable

class UnaryOp(Expression):
    """
    A unary operator expression within a statement
    """
    __slots__ = ('operator', 'right', )

    operator: str
    right: Expression

    def __init__(self, operator: str, right: Expression) -> None:
        super().__init__()

        self.operator = operator
        self.right = right

class BinaryOp(Expression):
    """
    A binary operator expression within a statement
    """
    __slots__ = ('operator', 'left', 'right', )

    operator: str
    left: Expression
    right: Expression

    def __init__(self, operator: str, left: Expression, right: Expression) -> None:
        super().__init__()

        self.operator = operator
        self.left = left
        self.right = right

class FunctionCall(Expression):
    """
    A function call expression within a statement
    """
    __slots__ = ('function', 'arguments', )

    function: 'Function'
    arguments: List[Expression]

    def __init__(self, function: 'Function') -> None:
        super().__init__()

        self.function = function
        self.arguments = []

class Subscript(Expression):
    """
    A subscript, for example to refer to a static array or tuple
    by index
    """
    __slots__ = ('varref', 'index', )

    varref: VariableReference
    index: Expression

    def __init__(self, varref: VariableReference, index: Expression) -> None:
        super().__init__()

        self.varref = varref
        self.index = index

class AccessStructMember(Expression):
    """
    Access a struct member for reading of writing
    """
    __slots__ = ('varref', 'struct_type3', 'member', )

    varref: VariableReference
    struct_type3: StructType3
    member: str

    def __init__(self, varref: VariableReference,  struct_type3: StructType3, member: str) -> None:
        super().__init__()

        self.varref = varref
        self.struct_type3 = struct_type3
        self.member = member

class Fold(Expression):
    """
    A (left or right) fold
    """
    class Direction(enum.Enum):
        """
        Which direction to fold in
        """
        LEFT = 0
        RIGHT = 1

    dir: Direction
    func: 'Function'
    base: Expression
    iter: Expression

    def __init__(
            self,
            dir_: Direction,
            func: 'Function',
            base: Expression,
            iter_: Expression,
        ) -> None:
        super().__init__()

        self.dir = dir_
        self.func = func
        self.base = base
        self.iter = iter_

class Statement:
    """
    A statement within a function
    """
    __slots__ = ()

class StatementPass(Statement):
    """
    A pass statement
    """
    __slots__ = ()

class StatementReturn(Statement):
    """
    A return statement within a function
    """
    __slots__ = ('value', )

    def __init__(self, value: Expression) -> None:
        self.value = value

class StatementIf(Statement):
    """
    An if statement within a function
    """
    __slots__ = ('test', 'statements', 'else_statements', )

    test: Expression
    statements: List[Statement]
    else_statements: List[Statement]

    def __init__(self, test: Expression) -> None:
        self.test = test
        self.statements = []
        self.else_statements = []

class FunctionParam:
    """
    A parameter for a Function
    """
    __slots__ = ('name', 'type3', 'type_str', )

    name: str
    type3: Type3
    type_str: str

    def __init__(self, name: str, type3: Type3) -> None:
        self.name = name
        self.type3 = type3
        self.type_str = type3.name

class Function:
    """
    A function processes input and produces output
    """
    __slots__ = ('name', 'lineno', 'exported', 'imported', 'statements', 'returns_type3', 'returns_str', 'posonlyargs', )

    name: str
    lineno: int
    exported: bool
    imported: bool
    statements: List[Statement]
    returns_type3: Type3
    returns_str: str
    posonlyargs: List[FunctionParam]

    def __init__(self, name: str, lineno: int) -> None:
        self.name = name
        self.lineno = lineno
        self.exported = False
        self.imported = False
        self.statements = []
        self.returns_type3 = type3types.none
        self.returns_str = 'None'
        self.posonlyargs = []

class StructDefinition:
    """
    The definition for a struct
    """
    __slots__ = ('struct_type3', 'lineno', )

    struct_type3: StructType3
    lineno: int

    def __init__(self, struct_type3: StructType3, lineno: int) -> None:
        self.struct_type3 = struct_type3
        self.lineno = lineno

class StructConstructor(Function):
    """
    The constructor method for a struct

    A function will generated to instantiate a struct. The arguments
    will be the defaults
    """
    __slots__ = ('struct_type3', )

    struct_type3: StructType3

    def __init__(self, struct_type3: StructType3) -> None:
        super().__init__(f'@{struct_type3.name}@__init___@', -1)

        self.returns_type3 = struct_type3

        for mem, typ in struct_type3.members.items():
            self.posonlyargs.append(FunctionParam(mem, typ, ))

        self.struct_type3 = struct_type3

# TODO: Broken after new type system
# class TupleConstructor(Function):
#     """
#     The constructor method for a tuple
#     """
#     __slots__ = ('tuple', )
#
#     tuple: TypeTuple
#
#     def __init__(self, tuple_: TypeTuple) -> None:
#         name = tuple_.render_internal_name()
#
#         super().__init__(f'@{name}@__init___@', -1)
#
#         self.returns = tuple_
#
#         for mem in tuple_.members:
#             self.posonlyargs.append(FunctionParam(f'arg{mem.idx}', mem.type, ))
#
#         self.tuple = tuple_

class ModuleConstantDef:
    """
    A constant definition within a module
    """
    __slots__ = ('name', 'lineno', 'type3', 'type_str', 'constant', 'data_block', )

    name: str
    lineno: int
    type3: Type3
    type_str: str
    constant: Constant
    data_block: Optional['ModuleDataBlock']

    def __init__(self, name: str, lineno: int, type3: Type3, constant: Constant, data_block: Optional['ModuleDataBlock']) -> None:
        self.name = name
        self.lineno = lineno
        self.type3 = type3
        self.type_str = type3.name
        self.constant = constant
        self.data_block = data_block

class ModuleDataBlock:
    """
    A single allocated block for module data
    """
    __slots__ = ('data', 'address', )

    data: List[ConstantPrimitive]
    address: Optional[int]

    def __init__(self, data: List[ConstantPrimitive]) -> None:
        self.data = data
        self.address = None

class ModuleData:
    """
    The data for when a module is loaded into memory
    """
    __slots__ = ('blocks', )

    blocks: List[ModuleDataBlock]

    def __init__(self) -> None:
        self.blocks = []

class Module:
    """
    A module is a file and consists of functions
    """
    __slots__ = ('data', 'types',  'struct_definitions', 'constant_defs', 'functions',)

    data: ModuleData
    struct_definitions: Dict[str, StructDefinition]
    constant_defs: Dict[str, ModuleConstantDef]
    functions: Dict[str, Function]

    def __init__(self) -> None:
        self.data = ModuleData()
        self.struct_definitions = {}
        self.constant_defs = {}
        self.functions = {}