"""
Contains the syntax tree for ourlang
"""
from __future__ import annotations

from typing import Dict, Iterable, List, Optional, Union

from .build.base import BuildBase
from .type5 import record as type5record
from .type5 import typeexpr as type5typeexpr


class SourceRef:
    __slots__ = ('filename', 'lineno', 'colno', )

    filename: str | None
    lineno: int | None
    colno: int | None

    def __init__(self, filename: str | None, lineno: int | None = None, colno: int | None = None) -> None:
        self.filename = filename
        self.lineno = lineno
        self.colno = colno

    def __repr__(self) -> str:
        return f"SourceRef({self.filename!r}, {self.lineno!r}, {self.colno!r})"

    def __str__(self) -> str:
        return f"{self.filename}:{self.lineno:>4}:{self.colno:<3}"

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

    sourceref: SourceRef
    type5: type5typeexpr.TypeExpr | None

    def __init__(self, *, sourceref: SourceRef) -> None:
        self.sourceref = sourceref
        self.type5 = None

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: int | float

    def __init__(self, value: int | float, sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)
        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', sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)
        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', sourceref: SourceRef) -> None:
        super().__init__(data_block, sourceref=sourceref)
        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', sourceref: SourceRef) -> None:
        super().__init__(data_block, sourceref=sourceref)
        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_type5', 'value', )

    struct_type5: type5record.Record
    value: List[Union[ConstantPrimitive, ConstantBytes, ConstantTuple, 'ConstantStruct']]

    def __init__(
            self,
            struct_type5: type5record.Record,
            value: List[Union[ConstantPrimitive, ConstantBytes, ConstantTuple, 'ConstantStruct']],
            data_block: 'ModuleDataBlock',
            sourceref: SourceRef
        ) -> None:
        super().__init__(data_block, sourceref=sourceref)
        self.struct_type5 = struct_type5
        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({self.struct_type5!r}, {self.value!r}, @{self.data_block.address!r})'

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'], sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)
        self.variable = variable

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

    operator: FunctionInstance
    left: Expression
    right: Expression

    def __init__(self, operator: FunctionInstance, left: Expression, right: Expression, sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)

        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 FunctionInstance(Expression):
    """
    When calling a polymorphic function with concrete arguments, we can generate
    code for that specific instance of the function.
    """
    __slots__ = ('function', )

    function: Union['Function', 'FunctionParam']

    def __init__(self, function: Union['Function', 'FunctionParam'], sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)

        self.function = function

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

    function_instance: FunctionInstance
    arguments: List[Expression]

    def __init__(self, function_instance: FunctionInstance, sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)

        self.function_instance = function_instance
        self.arguments = []

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

    function: 'Function'

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

class TupleInstantiation(Expression):
    """
    Instantiation a tuple
    """
    __slots__ = ('elements', )

    elements: List[Expression]

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

        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, sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)

        self.varref = varref
        self.index = index

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

    varref: VariableReference
    member: str

    def __init__(self, varref: VariableReference,  member: str, sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)

        self.varref = varref
        self.member = member

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

    sourceref: SourceRef

    def __init__(self, *, sourceref: SourceRef) -> None:
        self.sourceref = sourceref

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

    def __init__(self, sourceref: SourceRef) -> None:
        super().__init__(sourceref=sourceref)

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

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

        self.value = value

    def __repr__(self) -> str:
        return f'StatementReturn({repr(self.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', 'type5', )

    name: str
    type5: type5typeexpr.TypeExpr

    def __init__(self, name: str, type5: type5typeexpr.TypeExpr) -> None:
        assert type5typeexpr.is_concrete(type5)

        self.name = name
        self.type5 = type5

    def __repr__(self) -> str:
        return f'FunctionParam({self.name!r}, {self.type5!r})'

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

    name: str
    sourceref: SourceRef
    exported: bool
    imported: Optional[str]
    statements: List[Statement]
    type5: type5typeexpr.TypeExpr | type5typeexpr.ConstrainedExpr | None
    posonlyargs: List[FunctionParam]  # TODO: Replace me with arg names
    arg_names: list[str]

    def __init__(self, name: str, sourceref: SourceRef) -> None:
        self.name = name
        self.sourceref = sourceref
        self.exported = False
        self.imported = None
        self.statements = []
        self.type5 = None
        self.posonlyargs = []
        self.arg_names = []

class BuiltinFunction(Function):
    def __init__(self, name: str, type5: type5typeexpr.TypeExpr | type5typeexpr.ConstrainedExpr) -> None:
        super().__init__(name, SourceRef("/", 0, 0))
        self.type5 = type5

class StructDefinition:
    """
    The definition for a struct
    """
    __slots__ = ('struct_type5', 'sourceref', )

    struct_type5: type5record.Record
    sourceref: SourceRef

    def __init__(self, struct_type5: type5record.Record, sourceref: SourceRef) -> None:
        self.struct_type5 = struct_type5
        self.sourceref = sourceref

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

    struct_type5: type5record.Record

    def __init__(self, struct_type5: type5record.Record, sourceref: SourceRef) -> None:
        super().__init__(f'@{struct_type5.name}@__init___@', sourceref)
        self.struct_type5 = struct_type5

        for mem, typ in struct_type5.fields:
            self.arg_names.append(mem)
            self.posonlyargs.append(FunctionParam(mem, typ))

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

    name: str
    sourceref: SourceRef
    type5: type5typeexpr.TypeExpr
    constant: Constant

    def __init__(self, name: str, sourceref: SourceRef, type5: type5typeexpr.TypeExpr, constant: Constant) -> None:
        self.name = name
        self.sourceref = sourceref
        self.type5 = type5
        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

    def __repr__(self) -> str:
        return f'ModuleDataBlock({self.data!r}, {self.address!r})'

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[G]:
    """
    A module is a file and consists of functions
    """
    __slots__ = ('build', 'filename', 'data', 'types', 'type5s', 'struct_definitions', 'constant_defs', 'functions', 'methods', 'operators', 'functions_table', )

    build: BuildBase[G]
    filename: str
    data: ModuleData
    types: dict[str, type5typeexpr.TypeExpr]
    struct_definitions: Dict[str, StructDefinition]
    constant_defs: Dict[str, ModuleConstantDef]
    functions: Dict[str, Function]
    methods: Dict[str, type5typeexpr.TypeExpr | type5typeexpr.ConstrainedExpr]
    operators: Dict[str, type5typeexpr.TypeExpr | type5typeexpr.ConstrainedExpr]
    functions_table: dict[Function, int]

    def __init__(self, build: BuildBase[G], filename: str) -> None:
        self.build = build
        self.filename = filename

        self.data = ModuleData()
        self.types = {}
        self.struct_definitions = {}
        self.constant_defs = {}
        self.functions = {}
        self.methods = {}
        self.operators = {}
        self.functions_table = {}