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

import enum

from typing_extensions import Final

WEBASSEMBLY_BUILDIN_FLOAT_OPS: Final = ('abs', 'sqrt', 'ceil', 'floor', 'trunc', 'nearest', )
WEBASSEMBLY_BUILDIN_BYTES_OPS: Final = ('len', )

from .typing import (
    TypeBase,
    TypeNone,
    TypeBool,
    TypeUInt8, TypeUInt32, TypeUInt64,
    TypeInt32, TypeInt64,
    TypeFloat32, TypeFloat64,
    TypeBytes,
    TypeTuple, TypeTupleMember,
    TypeStaticArray, TypeStaticArrayMember,
    TypeStruct, TypeStructMember,

    TypeVar,
)

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

    type: TypeBase
    type_var: Optional[TypeVar]

    def __init__(self, type_: TypeBase) -> None:
        self.type = type_
        self.type_var = None

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

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

    value: int

    def __init__(self, type_: TypeUInt8, value: int) -> None:
        super().__init__(type_)
        self.value = value

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

    value: int

    def __init__(self, type_: TypeUInt32, value: int) -> None:
        super().__init__(type_)
        self.value = value

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

    value: int

    def __init__(self, type_: TypeUInt64, value: int) -> None:
        super().__init__(type_)
        self.value = value

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

    value: int

    def __init__(self, type_: TypeInt32, value: int) -> None:
        super().__init__(type_)
        self.value = value

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

    value: int

    def __init__(self, type_: TypeInt64, value: int) -> None:
        super().__init__(type_)
        self.value = value

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

    value: float

    def __init__(self, type_: TypeFloat32, value: float) -> None:
        super().__init__(type_)
        self.value = value

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

    value: float

    def __init__(self, type_: TypeFloat64, value: float) -> None:
        super().__init__(type_)
        self.value = value

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

    value: List[Constant]

    def __init__(self, type_: TypeTuple, value: List[Constant]) -> None:
        super().__init__(type_)
        self.value = value

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

    value: List[Constant]

    def __init__(self, type_: TypeStaticArray, value: List[Constant]) -> None:
        super().__init__(type_)
        self.value = value

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

    variable: 'FunctionParam' # also possibly local

    def __init__(self, type_: TypeBase, variable: 'FunctionParam') -> None:
        super().__init__(type_)
        self.variable = variable

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

    operator: str
    right: Expression

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

        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, type_: TypeBase, operator: str, left: Expression, right: Expression) -> None:
        super().__init__(type_)

        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__(function.returns)

        self.function = function
        self.arguments = []

class AccessBytesIndex(Expression):
    """
    Access a bytes index for reading
    """
    __slots__ = ('varref', 'index', )

    varref: VariableReference
    index: Expression

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

        self.varref = varref
        self.index = index

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

    varref: VariableReference
    member: TypeStructMember

    def __init__(self, varref: VariableReference,  member: TypeStructMember) -> None:
        super().__init__(member.type)

        self.varref = varref
        self.member = member

class AccessTupleMember(Expression):
    """
    Access a tuple member for reading of writing
    """
    __slots__ = ('varref', 'member', )

    varref: VariableReference
    member: TypeTupleMember

    def __init__(self, varref: VariableReference, member: TypeTupleMember, ) -> None:
        super().__init__(member.type)

        self.varref = varref
        self.member = member

class AccessStaticArrayMember(Expression):
    """
    Access a tuple member for reading of writing
    """
    __slots__ = ('varref', 'static_array', 'member', )

    varref: Union['ModuleConstantReference', VariableReference]
    static_array: TypeStaticArray
    member: Union[Expression, TypeStaticArrayMember]

    def __init__(self, varref: Union['ModuleConstantReference', VariableReference], static_array: TypeStaticArray, member: Union[TypeStaticArrayMember, Expression], ) -> None:
        super().__init__(static_array.member_type)

        self.varref = varref
        self.static_array = static_array
        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,
            type_: TypeBase,
            dir_: Direction,
            func: 'Function',
            base: Expression,
            iter_: Expression,
        ) -> None:
        super().__init__(type_)

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

class ModuleConstantReference(Expression):
    """
    An reference to a module constant expression within a statement
    """
    __slots__ = ('definition', )

    definition: 'ModuleConstantDef'

    def __init__(self, type_: TypeBase, definition: 'ModuleConstantDef') -> None:
        super().__init__(type_)
        self.definition = definition

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:
    __slots__ = ('name', 'type', 'type_var', )

    name: str
    type: TypeBase
    type_var: Optional[TypeVar]

    def __init__(self, name: str, type_: TypeBase) -> None:
        self.name = name
        self.type = type_
        self.type_var = None

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

    name: str
    lineno: int
    exported: bool
    imported: bool
    statements: List[Statement]
    returns: TypeBase
    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 = TypeNone()
        self.posonlyargs = []

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', )

    struct: TypeStruct

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

        self.returns = struct

        for mem in struct.members:
            self.posonlyargs.append(FunctionParam(mem.name, mem.type, ))

        self.struct = struct

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', 'type', 'constant', 'data_block', )

    name: str
    lineno: int
    type: TypeBase
    constant: Constant
    data_block: Optional['ModuleDataBlock']

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

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

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

    def __init__(self, data: List[Constant]) -> 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',  'structs', 'constant_defs', 'functions',)

    data: ModuleData
    types: Dict[str, TypeBase]
    structs: Dict[str, TypeStruct]
    constant_defs: Dict[str, ModuleConstantDef]
    functions: Dict[str, Function]

    def __init__(self) -> None:
        self.types = {
            'None': TypeNone(),
            'u8': TypeUInt8(),
            'u32': TypeUInt32(),
            'u64': TypeUInt64(),
            'i32': TypeInt32(),
            'i64': TypeInt64(),
            'f32': TypeFloat32(),
            'f64': TypeFloat64(),
            'bytes': TypeBytes(),
        }
        self.data = ModuleData()
        self.structs = {}
        self.constant_defs = {}
        self.functions = {}