import logging
from . import astnodes as ast
from .scope import SemanticError


class TypeChecker:
    """ Type checker """

    logger = logging.getLogger("c3check")

    def __init__(self, diag, context):
        self.diag = diag
        self.context = context

    def check(self):
        """ Check everything """
        for module in self.context.modules:
            self.check_module(module)

    def check_module(self, module: ast.Module):
        """ Check a module """
        assert isinstance(module, ast.Module)
        self.module_ok = True
        self.logger.info("Checking module %s", module.name)
        try:
            # Check defined types of this module:
            for typ in module.types:
                self.check_type(typ)

            for con in module.constants:
                self.check_type(con.typ)
                self.check_expr(con.value)
                con.value = self.do_coerce(con.value, con.typ)

            # Check global variables:
            for var in module.variables:
                assert not var.isLocal
                self.check_variable(var)
        except SemanticError as ex:
            self.error(ex.msg, ex.loc)

        # Check functions:
        for func in module.functions:
            # Try per function, in case of error, continue with next
            try:
                self.check_function(func)
            except SemanticError as ex:
                self.error(ex.msg, ex.loc)

        if not self.module_ok:
            raise SemanticError("Errors occurred", None)

    def check_type(self, typ, first=True, byname=False):
        """ Check a type.

        Determine struct offsets and check for recursiveness by using
        mark and sweep algorithm.

        The calling function could call this function with first set
        to clear the marks.
        """

        # Reset the mark and sweep:
        if first:
            self.got_types = set()

        # Resolve the type:
        typ = self.context.get_type(typ, not byname)

        # Check for recursion:
        if typ in self.got_types:
            raise SemanticError("Recursive data type {}".format(typ), None)

        if isinstance(typ, ast.BaseType):
            pass
        elif isinstance(typ, ast.PointerType):
            # If a pointed type is detected, stop structural
            # equivalence. This to allow linked lists!
            self.check_type(typ.ptype, first=False, byname=True)
        elif isinstance(typ, ast.StructureType):
            self.got_types.add(typ)
            # Setup offsets of fields. Is this the right place?:
            # TODO: move this struct offset calculation.
            offset = 0
            for struct_member in typ.fields:
                self.check_type(struct_member.typ, first=False)
                struct_member.offset = offset
                offset = offset + self.context.size_of(struct_member.typ)
        elif isinstance(typ, ast.ArrayType):
            self.check_type(typ.element_type, first=False)
        elif isinstance(typ, ast.DefinedType):
            pass
        else:  # pragma: no cover
            raise NotImplementedError("{} not implemented".format(type(typ)))

    def check_variable(self, var):
        """ Check a variable and especially its initial value """
        self.check_type(var.typ)
        if var.ival:
            var.ival = self.check_initial_value(var.ival, var.typ)

    def check_initial_value(self, ival, typ):
        """ Ensure that the initial value fits the given type """
        typ = self.context.get_type(typ)
        if isinstance(typ, ast.ArrayType):
            if not isinstance(ival, ast.ExpressionList):
                raise SemanticError("Invalid array initialization", ival.loc)

            # TODO: can we do legit constant evaluation here?
            size = self.context.eval_const(typ.size)
            if len(ival.expressions) != size:
                raise SemanticError(
                    "{} initial values given, expected {}".format(
                        len(ival.expressions), size
                    ),
                    ival.loc,
                )
            new_expressions = []
            for expr in ival.expressions:
                new_expr = self.check_initial_value(expr, typ.element_type)
                new_expressions.append(new_expr)
            ival.expressions = new_expressions
            return ival
        elif isinstance(typ, ast.StructureType):
            print(ival)
            if not isinstance(ival, ast.NamedExpressionList):
                raise SemanticError("Invalid struct initialization", ival.loc)
            if len(ival.expressions) != len(typ.fields):
                raise SemanticError("Wrong number of given fields", ival.loc)
            new_expressions = []
            for ival2, fld in zip(ival.expressions, typ.fields):
                field, expr = ival2
                if field != fld.name:
                    raise SemanticError("Wrong name of struct field", expr.loc)
                new_expr = self.check_initial_value(expr, fld.typ)
                new_expressions.append((field, new_expr))
            ival.expressions = new_expressions
            return ival
        else:
            self.check_expr(ival)
            return self.do_coerce(ival, typ)

    def check_function(self, function):
        """ Check a function. """
        for param in function.parameters:
            self.check_type(param.typ)

            # Parameters can only be simple types (pass by value)
            if not self.context.is_simple_type(param.typ):
                raise SemanticError(
                    "Function parameters can only be simple types",
                    function.loc,
                )

        if not self.context.is_simple_type(function.typ.returntype):
            raise SemanticError(
                "Functions can only return simple types", function.loc
            )

        for sym in function.inner_scope:
            self.check_type(sym.typ)

        self.current_function = function
        if function.body:
            self.check_stmt(function.body)
        self.current_function = None

    def check_stmt(self, code: ast.Statement):
        """ Check a statement """
        try:
            assert isinstance(code, ast.Statement)
            if isinstance(code, ast.Compound):
                for statement in code.statements:
                    self.check_stmt(statement)
            elif isinstance(code, ast.Empty):
                pass
            elif isinstance(code, ast.Assignment):
                self.check_assignment_stmt(code)
            elif isinstance(code, ast.VariableDeclaration):
                self.check_variable(code.var)
            elif isinstance(code, ast.ExpressionStatement):
                # Check that this is always a void function call
                if not isinstance(code.ex, ast.FunctionCall):
                    raise SemanticError("Not a call expression", code.ex.loc)
                value = self.check_function_call(code.ex)
                if not self.context.equal_types("void", code.ex.typ):
                    raise SemanticError(
                        "Can only call void functions", code.ex.loc
                    )
                assert value is None
            elif isinstance(code, ast.If):
                self.check_if_stmt(code)
            elif isinstance(code, ast.Return):
                self.check_return_stmt(code)
            elif isinstance(code, ast.While):
                self.check_while(code)
            elif isinstance(code, ast.For):
                self.check_for_stmt(code)
            elif isinstance(code, ast.Switch):
                self.check_switch_stmt(code)
            else:  # pragma: no cover
                raise NotImplementedError(str(code))
        except SemanticError as exc:
            self.error(exc.msg, exc.loc)

    def check_if_stmt(self, code):
        """ Check if statement """