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

from .build.base import BuildBase
from .type3.functions import FunctionSignature, TypeVariableContext
from .type3.typeclasses import Type3ClassMethod
from .type3.types import Type3, TypeApplication_Struct
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__ = ('type3', 'type5', 'sourceref', )

    sourceref: SourceRef | None
    type3: Type3 | None
    type5: type5typeexpr.TypeExpr | None

    def __init__(self, *, sourceref: SourceRef | None = None) -> None:
        self.sourceref = sourceref
        self.type3 = None
        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: Union[int, float]

    def __init__(self, value: Union[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') -> 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_type3', 'value', )

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

    def __init__(self, struct_type3: Type3, value: List[Union[ConstantPrimitive, ConstantBytes, ConstantTuple, 'ConstantStruct']], data_block: 'ModuleDataBlock') -> None:
        super().__init__(data_block)
        self.struct_type3 = struct_type3
        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_type3!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: Type3ClassMethod
    left: Expression
    right: Expression

    def __init__(self, operator: Type3ClassMethod, 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 FunctionCall(Expression):
    """
    A function call expression within a statement
    """
    __slots__ = ('function', 'arguments', )

    function: Union['Function', 'FunctionParam', Type3ClassMethod]
    arguments: List[Expression]

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

        self.function = function
        self.arguments = []

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

    function: 'Function'

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

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: Type3
    member: str

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

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

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

    sourceref: SourceRef | None

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

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

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

    name: str
    type3: Type3
    type5: type5typeexpr.TypeExpr

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

        self.name = name
        self.type3 = type3
        self.type5 = type5

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

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

    name: str
    sourceref: SourceRef
    exported: bool
    imported: Optional[str]
    statements: List[Statement]
    signature: FunctionSignature
    returns_type3: Type3
    type5: type5typeexpr.TypeExpr | None
    posonlyargs: List[FunctionParam]

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

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

    struct_type3: Type3
    sourceref: SourceRef

    def __init__(self, struct_type3: Type3, sourceref: SourceRef) -> None:
        self.struct_type3 = struct_type3
        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_type3', )

    struct_type3: Type3

    def __init__(self, struct_type3: Type3) -> None:
        super().__init__(f'@{struct_type3.name}@__init___@', -1, struct_type3)  # type: ignore # TODO

        assert isinstance(struct_type3.application, TypeApplication_Struct)

        for mem, typ in struct_type3.application.arguments:
            self.posonlyargs.append(FunctionParam(mem, typ))  # type: ignore # TODO
            self.signature.args.append(typ)

        self.signature.args.append(struct_type3)

        self.struct_type3 = struct_type3

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

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

    def __init__(self, name: str, sourceref: SourceRef, type3: Type3, type5: type5typeexpr.TypeExpr, constant: Constant) -> None:
        self.name = name
        self.sourceref = sourceref
        self.type3 = type3
        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, Type3]
    type5s: dict[str, type5typeexpr.TypeExpr]
    struct_definitions: Dict[str, StructDefinition]
    constant_defs: Dict[str, ModuleConstantDef]
    functions: Dict[str, Function]
    methods: Dict[str, Type3ClassMethod]
    operators: Dict[str, Type3ClassMethod]
    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.type5s = {}
        self.struct_definitions = {}
        self.constant_defs = {}
        self.functions = {}
        self.methods = {}
        self.operators = {}
        self.functions_table = {}