# Visitor to *typecheck* MiniC files
from typing import List
from MiniCVisitor import MiniCVisitor
from MiniCParser import MiniCParser
from Lib.Errors import MiniCInternalError, MiniCTypeError

from enum import Enum


class BaseType(Enum):
    Float, Integer, Boolean, String = range(4)


# Basic Type Checking for MiniC programs.
class MiniCTypingVisitor(MiniCVisitor):

    def __init__(self):
        self._memorytypes = dict()  # id -> types
        # For now, we don't have real functions ...
        self._current_function = "main"

    def _raise(self, ctx, for_what, *types):
        raise MiniCTypeError(
            'In function {}: Line {} col {}: invalid type for {}: {}'.format(
                self._current_function,
                ctx.start.line, ctx.start.column, for_what,
                ' and '.join(t.name.lower() for t in types)))

    def _assertSameType(self, ctx, for_what, *types):
        if not all(types[0] == t for t in types):
            raise MiniCTypeError(
                'In function {}: Line {} col {}: type mismatch for {}: {}'.format(
                    self._current_function,
                    ctx.start.line, ctx.start.column, for_what,
                    ' and '.join(t.name.lower() for t in types)))

    def _raiseNonType(self, ctx, message):
        raise MiniCTypeError(
            'In function {}: Line {} col {}: {}'.format(
                self._current_function,
                ctx.start.line, ctx.start.column, message))

    # type declaration

    def visitVarDecl(self, ctx) -> None:
        ttype = self.visit(ctx.typee())
        ids = self.visit(ctx.id_l())
        for i in ids:
            if i in self._memorytypes:
                self._raiseNonType(ctx,"Variable {} already declared".format(i))
            self._memorytypes[i] = ttype

    def visitBasicType(self, ctx):
        assert ctx.mytype is not None
        if ctx.mytype.type == MiniCParser.INTTYPE:
            return BaseType.Integer
        elif ctx.mytype.type == MiniCParser.FLOATTYPE:
            return BaseType.Float
        if ctx.mytype.type == MiniCParser.STRINGTYPE:
            return BaseType.String
        elif ctx.mytype.type == MiniCParser.BOOLTYPE:
            return BaseType.Boolean
        else:
            raise MiniCInternalError("Unknown type '%s'" % ctx.mytype)

    def visitIdList(self, ctx) -> List[str]:
        l = self.visit(ctx.id_l())
        l.append(ctx.ID().getText())
        return l

    def visitIdListBase(self, ctx) -> List[str]:
        return [ctx.ID().getText()]

    # typing visitors for expressions, statements !

    # visitors for atoms --> type
    def visitParExpr(self, ctx):
        return self.visit(ctx.expr())

    def visitIntAtom(self, ctx):
        return BaseType.Integer

    def visitFloatAtom(self, ctx):
        return BaseType.Float

    def visitBooleanAtom(self, ctx):
        return BaseType.Boolean

    def visitIdAtom(self, ctx):
        try:
            return self._memorytypes[ctx.getText()]
        except KeyError:
            self._raiseNonType(ctx,
                               "Undefined variable {}".format(ctx.getText()))

    def visitStringAtom(self, ctx):
        return BaseType.String

    # now visit expr

    def visitAtomExpr(self, ctx):
        return self.visit(ctx.atom())

    def visitOrExpr(self, ctx):
        fstT = self.visit(ctx.expr(0))
        sndT = self.visit(ctx.expr(1))
        if(fstT == sndT):
            if(fstT == BaseType.Boolean):
                return BaseType.Boolean
            else:
                self._raise(ctx,"boolean operator", fstT)
        else:
            self._raise(ctx,"boolean operator", fstT,sndT)

    def visitAndExpr(self, ctx):
        fstT = self.visit(ctx.expr(0))
        sndT = self.visit(ctx.expr(1))
        if(fstT == sndT):
            if(fstT == BaseType.Boolean):
                return BaseType.Boolean
            else:
                self._raise(ctx,"boolean operator", fstT)
        else:
            self._raise(ctx,"boolean operator", fstT,sndT)

    def visitEqualityExpr(self, ctx):
        fstT = self.visit(ctx.expr(0))
        sndT = self.visit(ctx.expr(1))
        if(fstT == sndT):
            return BaseType.Boolean
        else:
            self._assertSameType(ctx,"equality operator", fstT,sndT)

    def visitRelationalExpr(self, ctx):
        fstT = self.visit(ctx.expr(0))
        sndT = self.visit(ctx.expr(1))
        if(fstT == sndT):
            if(fstT in [BaseType.Float, BaseType.Integer]):
                return BaseType.Boolean
            else:
                self._raise(ctx, "comparaison operator", fstT)
        else:
            self._raise(ctx,"comparaison operator", fstT,sndT)

    def visitAdditiveExpr(self, ctx):
        assert ctx.myop is not None
        fstT = self.visit(ctx.expr(0))
        sndT = self.visit(ctx.expr(1))
        if(fstT == sndT):
            if(fstT in [BaseType.Float, BaseType.Integer]):
                return fstT
            elif (fstT == BaseType.String and ctx.myop.type == MiniCParser.PLUS):
                return BaseType.String
            else:
                self._raise(ctx, "additive operands", fstT,sndT)
        else:
            self._raise(ctx,"additive operator", fstT,sndT)

    def visitMultiplicativeExpr(self, ctx):
        fstT = self.visit(ctx.expr(0))
        sndT = self.visit(ctx.expr(1))
        if(fstT == sndT):
            if(not fstT in [BaseType.Float, BaseType.Integer]):
                self._raise(ctx, "multiplicative operands", fstT)
            else:
                return fstT
        else:
            self._raise(ctx,"multiplicative operands", fstT,sndT)

    def visitNotExpr(self, ctx):
        fstT = self.visit(ctx.expr())
        if(fstT == BaseType.Boolean):
            return BaseType.Boolean
        else:
            self._raise(ctx,"not operator", fstT)

    def visitUnaryMinusExpr(self, ctx):
        fstT = self.visit(ctx.expr())
        if(fstT in [BaseType.Integer,BaseType.Float]):
            return fstT
        else:
            self._raise(ctx,"minus operator", fstT)

    # visit statements

    def visitPrintlnintStat(self, ctx):
        etype = self.visit(ctx.expr())
        if etype != BaseType.Integer:
            self._raise(ctx, 'println_int statement', etype)

    def visitPrintlnfloatStat(self, ctx):
        etype = self.visit(ctx.expr())
        if etype != BaseType.Float:
            self._raise(ctx, 'println_float statement', etype)

    def visitPrintlnboolStat(self, ctx):
        etype = self.visit(ctx.expr())
        if etype != BaseType.Boolean:
            self._raise(ctx, 'println_bool statement', etype)

    def visitPrintlnstringStat(self, ctx):
        etype = self.visit(ctx.expr())
        if etype != BaseType.String:
            self._raise(ctx, 'println_string statement', etype)

    def visitAssignStat(self, ctx):
        theid = ctx.ID().getText()
        ttyp = self.visit(ctx.expr())
        if(not theid in self._memorytypes):
            self._raiseNonType(ctx,"Undefined variable {}".format(theid))
        if(self._memorytypes[theid] != ttyp):
            self._assertSameType(ctx,theid, self._memorytypes[theid],ttyp)

    def visitWhileStat(self, ctx):
        cond = self.visit(ctx.expr())
        if(cond != BaseType.Boolean):
            self._raise(ctx, "while condition", cond)
        self.visit(ctx.body)

    def visitIfStat(self, ctx):
        cond = self.visit(ctx.expr())
        if(cond != BaseType.Boolean):
            self._raise(ctx, "if condition", cond)
        self.visit(ctx.then_block)
        if(ctx.else_block != None):
            self.visit(ctx.else_block)
    
    def visitForStat(self, ctx):
        if(ctx.init_stat != None):
            self.visit(ctx.init_stat)
        if(ctx.cond != None):
            cond = self.visit(ctx.cond)
            if(cond != BaseType.Boolean):
                self._raise(ctx, "for condition", cond)
        if(ctx.loop_stat != None):
            self.visit(ctx.loop_stat)
        
        self.visit(ctx.stat_block())