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

from typing_extensions import Final

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

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

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 ConstantMemoryStored(Constant):
    """
    An constant value expression within a statement

    # FIXME: Rename to literal
    """
    __slots__ = ('data_block', )

    data_block: 'ModuleDataBlock'

    def __init__(self, data_block: 'ModuleDataBlock') -> None:
        super().__init__()
        self.data_block = data_block

class ConstantBytes(ConstantMemoryStored):
    """
    A bytes constant value expression within a statement
    """
    __slots__ = ('value', )

    value: bytes

    def __init__(self, value: bytes, data_block: 'ModuleDataBlock') -> None:
        super().__init__(data_block)
        self.value = value

    def __repr__(self) -> str:
        # Do not repr the whole ModuleDataBlock
        # As this has a reference back to this constant for its data
        # which it needs to compile the data into the program
        return f'ConstantBytes({repr(self.value)}, @{repr(self.data_block.address)})'

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

    value: List[Union[ConstantPrimitive, ConstantBytes, 'ConstantTuple', 'ConstantStruct']]

    def __init__(self, value: List[Union[ConstantPrimitive, ConstantBytes, 'ConstantTuple', 'ConstantStruct']], data_block: 'ModuleDataBlock') -> None:
        super().__init__(data_block)
        self.value = value

    def __repr__(self) -> str:
        # Do not repr the whole ModuleDataBlock
        # As this has a reference back to this constant for its data
        # which it needs to compile the data into the program
        return f'ConstantTuple({repr(self.value)}, @{repr(self.data_block.address)})'

class ConstantStruct(ConstantMemoryStored):
    """
    A Struct constant value expression within a statement
    """
    __slots__ = ('struct_name', 'value', )

    struct_name: str
    value: List[Union[ConstantPrimitive, ConstantBytes, ConstantTuple, 'ConstantStruct']]

    def __init__(self, struct_name: str, value: List[Union[ConstantPrimitive, ConstantBytes, ConstantTuple, 'ConstantStruct']], data_block: 'ModuleDataBlock') -> None:
        super().__init__(data_block)
        self.struct_name = struct_name
        self.value = value

    def __repr__(self) -> str:
        # Do not repr the whole ModuleDataBlock
        # As this has a reference back to this constant for its data
        # which it needs to compile the data into the program
        return f'ConstantStruct({repr(self.struct_name)}, {repr(self.value)}, @{repr(self.data_block.address)})'

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

    def __repr__(self) -> str:
        return f'BinaryOp({repr(self.operator)}, {repr(self.left)}, {repr(self.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 TupleInstantiation(Expression):
    """
    Instantiation a tuple
    """
    __slots__ = ('elements', )

    elements: List[Expression]

    def __init__(self, elements: List[Expression]) -> None:
        super().__init__()

        self.elements = elements

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

    name: str
    type3: Type3OrPlaceholder

    def __init__(self, name: str, type3: Optional[Type3]) -> None:
        self.name = name
        self.type3 = PlaceholderForType([self]) if type3 is None else type3

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

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

    def __init__(self, name: str, lineno: int) -> None:
        self.name = name
        self.lineno = lineno
        self.exported = False
        self.imported = None
        self.statements = []
        self.returns_type3 = type3types.none # FIXME: This could be a placeholder
        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

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

    name: str
    lineno: int
    type3: Type3
    constant: Constant

    def __init__(self, name: str, lineno: int, type3: Type3, constant: Constant) -> None:
        self.name = name
        self.lineno = lineno
        self.type3 = type3
        self.constant = constant

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

    data: List[Union[ConstantPrimitive, ConstantMemoryStored]]
    address: Optional[int]

    def __init__(self, data: Iterable[Union[ConstantPrimitive, ConstantMemoryStored]]) -> 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 = {}