package chocopy.pa2;

import chocopy.common.analysis.AbstractNodeAnalyzer;
import chocopy.common.analysis.SymbolTable;
import chocopy.common.analysis.types.ClassVType;
import chocopy.common.analysis.types.ClassValueType;
import chocopy.common.analysis.types.FuncType;
import chocopy.common.analysis.types.FuncValueType;
import chocopy.common.analysis.types.ListValueType;
import chocopy.common.analysis.types.Type;
import chocopy.common.analysis.types.ValueType;
import chocopy.common.astnodes.*;

import static chocopy.common.analysis.types.Type.INT_TYPE;
import static chocopy.common.analysis.types.Type.OBJECT_TYPE;

import java.util.ArrayList;
import java.util.HashMap;

import javax.swing.text.StyledEditorKit.BoldAction;

import com.fasterxml.jackson.annotation.JacksonInject.Value;

/**
 * Analyzer that performs ChocoPy type checks on all nodes. Applied after collecting declarations.
 */
public class TypeChecker extends AbstractNodeAnalyzer<Type> {
    // global scope
    private final SymbolTable<Type> sym; 
    private final DeclarationAnalyzer declAnalyzer;
    /** The current symbol table (changes depending on the function being analyzed). */
    private SymbolTable<Type> currentScope;
    private Type currReturnType;
    private boolean returned = false, member = false;
    /** Collector for errors. */
    private final Errors errors;
    private boolean assign = false;
    private boolean declAnalyzed = false;
    private final HashMap<FuncDef, SymbolTable<Type>> funcScopes;
    /**
     * Creates a type checker using GLOBALSYMBOLS for the initial global symbol table and ERRORS0 to
     * receive semantic errors.
     */
    public TypeChecker(SymbolTable<Type> globalSymbols, Errors errors0) {
        sym = globalSymbols;
        currentScope = sym;
        errors = errors0;
        currReturnType = null;
        declAnalyzer = new DeclarationAnalyzer(errors0, this, globalSymbols);
        funcScopes = new HashMap<>();
    }
    /**
     * Inserts an error message in NODE if there isn't one already. The message is constructed with
     * MESSAGE and ARGS as for String.format.
     */
    private boolean isVariableType(Type ty){
        return ty.isSpecialType() || ty.equals(Type.OBJECT_TYPE);
    }
    public boolean pushDeclAnalyzed() {
        boolean orig = declAnalyzed;
        declAnalyzed = true;
        return orig;
    }
    public void popDeclAnalyzed(boolean orig) {
        declAnalyzed = orig;
    }
    private Type declAnalyze(Node node){
        //if(currentScope != sym)
        declAnalyzer.setScope(currentScope);
        return node.dispatch(declAnalyzer);
    }
    private void err(Node node, String message, Object... args) {
        errors.semError(node, message, args);
    }

    @Override
    public Type analyze(Program program) {
        for (Declaration decl : program.declarations) {
            decl.dispatch(this);
        }
        for (Stmt stmt : program.statements) {
            stmt.dispatch(this);
        }
        return null;
    }
    @Override
    public Type analyze(ClassDef node){
        ClassVType t = (ClassVType) sym.get(node.name.name);
        SymbolTable<Type> backScope = currentScope;
        currentScope = t.scope;
        declAnalyzer.setCurrClass(t);
        for(Declaration decl : node.declarations){
            decl.dispatch(this);
        }
        declAnalyzer.setCurrClass(null);
        currentScope = backScope;
        return null;
    }
    @Override
    public Type analyze(AssignStmt node) {
        Type tr = node.value.dispatch(this);
        Type tl;
        boolean error = false;
        assign=true;
        for (Expr ex : node.targets)
        {
            tl = ex.dispatch(this);
            if(error) continue;
            else if(tl == null)
            {
                err(node, "Expression `%s` type inference error.", ex);
                error = true;
            } 
            else if (ex instanceof IndexExpr && 
                ((IndexExpr)ex).list.getInferredType().equals(Type.STR_TYPE))
            {
                err(ex, "`str` is not a list type");
                error = true;
            }
            else if(tr!=null && tl.isListType() && tr.isListType())
            {
                if(!((!tl.elementType().isSpecialType()&&tr.elementType().equals(Type.NONE_TYPE))||
                    tl.equals(tr)||tr.elementType().equals(Type.EMPTY_TYPE)))
                    {
                        err(node, "Expected type `%s`; got type `%s`", tl, tr);
                    error = true;
                    }
            }
            else if(tl.isListType() && Type.EMPTY_TYPE.equals(tr)) ; //continue;
            else if(tr != null && !(StudentAnalysis.subClassOf(tl, tr, currentScope) || !tl.isSpecialType() && tr.equals(Type.NONE_TYPE)))
            {
                err(node, "Expected type `%s`; got type `%s`", tl, tr);
                error = true;
            }
        }
        assign=false;
        return null;
    }

    @Override
    public Type analyze(BooleanLiteral node) {
        return node.setInferredType(Type.BOOL_TYPE);
    }
    public void dispatchFuncDef(FuncDef node, SymbolTable<Type> scope, boolean declAnalyzed){
        boolean prevDeclAnalyzed = this.declAnalyzed;
        this.declAnalyzed = declAnalyzed;
        SymbolTable<Type> origScope = currentScope;
        currentScope = scope;
        node.dispatch(this);
        currentScope = origScope;
        this.declAnalyzed = prevDeclAnalyzed;
    }
    @Override
    public Type analyze(FuncDef node) {
        SymbolTable<Type> origScope = currentScope;
        if(funcScopes.get(node) != null)
            System.out.println("error");
        {
            currentScope = declAnalyzer.createScope(currentScope);
            funcScopes.put(node, currentScope);
            declAnalyzer.setPostCheck();
            declAnalyze(node);
            //currentScope = funcScopes.get(node);
            returned = false;
            Type prevReturnType = this.currReturnType;
            this.currReturnType = ValueType.annotationToValueType(node.returnType);
            for(Stmt st : node.statements)
                st.dispatch(this);

            if(currReturnType != null && currReturnType.isSpecialType() && !returned)
                err(node.name, "All paths in this function/method must have a return statement: %s", node.name.name);
            this.currReturnType = prevReturnType;
        }
        currentScope = origScope;
        return null;
    }

    @Override
    public Type analyze(CallExpr node) {
        Type f = currentScope.get(node.function.name);
        ArrayList<Type> types = new ArrayList<>();
        for(Expr ex: node.args)
            types.add(ex.dispatch(this));
        if(f != null && f.isFuncType())
        {
            FuncType fty = (FuncType) f;
            int lArgs = node.args.size(), lPars = fty.parameters.size();
            if(lArgs != lPars)
                err(node, "Expected %d arguments; got %d", lPars, lArgs);
            else{
                for(int i = 0; i < lArgs; ++i){
                    Type p = fty.parameters.get(i);
                    Type c = types.get(i);
                    if(((p.isSpecialType()&&!p.equals(c)) || 
                        (!p.isSpecialType()&&!StudentAnalysis.subClassOf(p, c, currentScope)))
                        &&!(p.isListType()&&c.equals(Type.EMPTY_TYPE))
                        )
                        err(node,"Expected type `%s`; got type `%s` in parameter %d", p, c, i);
                }
            }
            node.function.setInferredType(new FuncType(fty.parameters, fty.returnType));
            return node.setInferredType(fty.returnType);
        }
        else if (f != null && f instanceof ClassVType){
            ClassVType cty = (ClassVType) f;
            return node.setInferredType(new ClassValueType(f.className()));
        }
        else{
            err(node, "Not a function or class: %s", node.function.name);
            return node.setInferredType(Type.NONE_TYPE);
        }
    }

    @Override
    public Type analyze(ClassType node) {
        return sym.get(node.className);
    }

    @Override
    public Type analyze(ForStmt node) {
        Type iterableType = node.iterable.setInferredType(
                node.iterable.dispatch(this));
        if(iterableType == null)
            err(node, "Iterable `%s` type inference error.", node.iterable);
        else if (iterableType.equals(Type.STR_TYPE))
            node.identifier.setInferredType(Type.STR_TYPE);   
        else if(iterableType.elementType() == null){
            err(node, "`%s` isn't iterable", iterableType);
        }
        else
            node.identifier.setInferredType(
                iterableType.elementType()
            );    
        for(Stmt st : node.body)
            st.dispatch(this);
        return null;
    }

    @Override
    public Type analyze(IfExpr node) {
        Type condTy = node.condition.dispatch(this);
        if(!condTy.equals(Type.BOOL_TYPE)){
            err(node, "Condition expression cannot be of type `%s`", condTy.className());
        }
        
        Type ifTy = node.thenExpr.dispatch(this), 
             elseTy = node.elseExpr.dispatch(this);
        
        
        return node.setInferredType(StudentAnalysis.lowestCommonType(ifTy, elseTy, currentScope));
    }

    @Override
    public Type analyze(IfStmt node) {
        Type condTy = node.condition.dispatch(this);
        if(!condTy.equals(Type.BOOL_TYPE)){
            err(node, "Condition expression cannot be of type `%s`", condTy.className());
        }
        boolean prevReturned = returned, thenReturned;
        for(Stmt st : node.thenBody)
            st.dispatch(this);
        thenReturned = prevReturned || returned;
        returned = prevReturned;
        for(Stmt st : node.elseBody)
            st.dispatch(this);
        returned = returned && thenReturned;
        return null;
    }

    @Override
    public Type analyze(IndexExpr node) {
        Type listTy = node.list.dispatch(this);
        if(!(listTy.isListType() || listTy.equals(Type.STR_TYPE)))
            err(node, "Cannot index into type `%s`", listTy);
        if(!node.index.dispatch(this).equals(Type.INT_TYPE))
            err(node, "Index is of non-integer type `%s`", node.index.getInferredType());
        if(listTy.equals(Type.STR_TYPE))
            return node.setInferredType(Type.STR_TYPE);
        else if(listTy.elementType() != null)
            return node.setInferredType(listTy.elementType());
        else return node.setInferredType(Type.OBJECT_TYPE);
    }

    @Override
    public Type analyze(ListExpr node) {
        Type t = null;
        for(Expr ex : node.elements)
        {
            Type thisType = ex.dispatch(this);
            if(t == null)
                t = thisType;
            t = StudentAnalysis.lowestCommonType(t, thisType, currentScope);
        }
        if(t == null)
            return node.setInferredType(Type.EMPTY_TYPE);
        return node.setInferredType(new ListValueType(t));
    }

    @Override
    public Type analyze(MemberExpr node) {
        boolean prevIsMember = member;
        member = false;
        Type ty = node.object.dispatch(this);
        if(ty instanceof ClassValueType){
            ty = currentScope.get(((ClassValueType) ty).className());
            if(ty instanceof ClassVType){
                ClassVType classTy = (ClassVType) ty;
                Type type = classTy.scope == null? null:classTy.scope.get(node.member.name);
                if(type != null)
                    return node.setInferredType(type);
                else
                    err(node, "There is no %s named `%s` in class `%s`",
                    prevIsMember?"method":"attribute", node.member.name, classTy);
            } else 
                err(node, "Class `%s` undefined", ty.className());
        }
        else
            err(node, "`%s` isn't a class.", ty);
        return node.setInferredType(ValueType.OBJECT_TYPE);
    }

    @Override
    public Type analyze(MethodCallExpr node) {
        boolean prevIsMember = member;
        member = true;
        Type ty = node.method.dispatch(this);
        member = prevIsMember;
        Type thisTy = Type.OBJECT_TYPE;
        if(ty instanceof FuncType){
            FuncType funcTy = (FuncType) ty;
            int len = funcTy.parameters.size() - 1, largs = node.args.size();
            if(largs != len)
                err(node, "Expected %d arguments; got %d", len, largs);
            len = len<=largs?len:largs;
            for(int i = 0; i < len; ++i){
                Expr thisArg = node.args.get(i);
                Type thisArgTy = thisArg.setInferredType(thisArg.dispatch(this)), 
                     thisParamTy = funcTy.parameters.get(i + 1);

                if(!thisParamTy.equals(thisArgTy) && !StudentAnalysis.subClassOf(thisParamTy, thisArgTy, currentScope)
                &&!(thisParamTy.isListType()&&thisArgTy.equals(Type.EMPTY_TYPE)))
                    err(node, "Expected type `%s`; got type `%s` in parameter %d", 
                            thisParamTy, thisArgTy, i + 1);
            }
            thisTy = funcTy.returnType;
        }
        // else
        //     err(node.method.member, "`%s` isn't a MemberFunction.", node.method.member.name);

        for(Expr args : node.args)
            args.dispatch(this);

        return node.setInferredType(thisTy);
    }

    @Override
    public Type analyze(NoneLiteral node) {
        return node.setInferredType(Type.NONE_TYPE);
    }
    /*
    @Override
    public Type analyze(NonLocalDecl node) {
        SymbolTable<Type> parent = currentScope.getParent();
        if(parent == null || parent == sym || 
            !parent.getDeclaredSymbols().contains(node.variable.name)||
            !isVariableType(parent.get(node.variable.name))
        ){
                err(node.variable, "Not a nonlocal variable: %s", node.variable.name);
        } else if(currentScope.getDeclaredSymbols().contains(node.variable.name)){
            errors.semError(
                    node.variable, "Duplicate declaration of identifier in same scope: %s", node.variable.name);
        }
        else
        {
            Type nonlocalVar = parent.get(node.variable.name);
            currentScope.put(node.variable.name, nonlocalVar);
        }
        return null;
    }
    */
    @Override
    public Type analyze(ReturnStmt node) {
        if(node.value != null)
            node.value.dispatch(this);
        Type p = this.currReturnType;
        Type c = (node.value == null ? Type.NONE_TYPE: node.value.getInferredType());
        if(node.value == null && p.isSpecialType())
            err(node, "Expected type `%s`; got `None`", p);
        else if(p.isSpecialType()&&(c.equals(Type.NONE_TYPE)) || (!c.equals(Type.NONE_TYPE) && !StudentAnalysis.subClassOf(p, c, currentScope)))
            err(node, "Expected type `%s`; got type `%s`", p, c);
        returned = true;
        return null;
    }

    @Override
    public Type analyze(StringLiteral node) {
        return node.setInferredType(Type.STR_TYPE);
    }

    @Override
    public Type analyze(TypedVar node) {
        return ValueType.annotationToValueType(node.type);
    }

    @Override
    public Type analyze(VarDef node) {
        return declAnalyze(node);
    }

    @Override
    public Type analyze(UnaryExpr node) {
        Type t = node.operand.dispatch(this);
        switch (node.operator) {
            case "-":
            case "+":
                if (INT_TYPE.equals(t)) {
                    return node.setInferredType(INT_TYPE);
                } else {
                    err(node, "Cannot apply operator `%s` on type `%s`", node.operator, t);
                    return node.setInferredType(INT_TYPE);
                }
            case "not":
            if (!(Type.BOOL_TYPE.equals(t)))
                err(node, "Cannot apply operator `not` on type `%s`", t);
            return node.setInferredType(Type.BOOL_TYPE);
            default:
                return node.setInferredType(OBJECT_TYPE);
        }
    }


    @Override
    public Type analyze(WhileStmt node) {
        if(!node.condition.dispatch(this).equals(Type.BOOL_TYPE))
            err(node, "`%s` isn't a boolean expression.", node.condition);
        for(Stmt st : node.body)
            st.dispatch(this);
        return null;
    }


    @Override
    public Type analyze(ExprStmt s) {
        s.expr.dispatch(this);
        return null;
    }

    @Override
    public Type analyze(IntegerLiteral i) {
        return i.setInferredType(Type.INT_TYPE);
    }

    @Override
    public Type analyze(BinaryExpr e) {
        Type t1 = e.left.dispatch(this);
        Type t2 = e.right.dispatch(this);

        switch (e.operator) {
            case "-":
            case "*":
            case "//":
            case "%":
                if (INT_TYPE.equals(t1) && INT_TYPE.equals(t2)) {
                    return e.setInferredType(INT_TYPE);
                } else {
                    err(e, "Cannot apply operator `%s` on types `%s` and `%s`", e.operator, t1, t2);
                    return e.setInferredType(INT_TYPE);
                }
            case "+":
                if (INT_TYPE.equals(t1) && INT_TYPE.equals(t2)) {
                    return e.setInferredType(INT_TYPE);
                } else if(Type.STR_TYPE.equals(t1) && Type.STR_TYPE.equals(t2)){
                    return e.setInferredType(Type.STR_TYPE);
                } else if (t1.isListType() && t2.isListType()){
                    return e.setInferredType(new ListValueType(StudentAnalysis.
                        lowestCommonType(t1.elementType(), t2.elementType(), currentScope)));
                }else if (t1.isListType() && t2.equals(Type.EMPTY_TYPE))
                    return e.setInferredType(t1);
                else {
                    err(e, "Cannot apply operator `+` on types `%s` and `%s`", t1, t2); 
                    return e.setInferredType(INT_TYPE);
                }
            case "and":
            case "or":
            if (!(Type.BOOL_TYPE.equals(t1) && Type.BOOL_TYPE.equals(t2)))
                err(e, "Cannot apply operator `%s` on types `%s` and `%s`", e.operator, t1, t2);
            return e.setInferredType(Type.BOOL_TYPE);
            case ">":
            case "<":
            case ">=":
            case "<=":
            if (!(INT_TYPE.equals(t1) && INT_TYPE.equals(t2)))
                err(e, "Cannot apply operator `%s` on types `%s` and `%s`", e.operator, t1, t2);
            return e.setInferredType(Type.BOOL_TYPE);
            case "!=":
            case "==":
            if (!(INT_TYPE.equals(t1) && INT_TYPE.equals(t2) 
                || Type.BOOL_TYPE.equals(t1) && Type.BOOL_TYPE.equals(t2)
                || Type.STR_TYPE.equals(t1) && Type.STR_TYPE.equals(t2)))
                    err(e, "Cannot apply operator `%s` on types `%s` and `%s`", e.operator, t1, t2);
               return e.setInferredType(Type.BOOL_TYPE);
            case "is":
            if(t1.isSpecialType()||t2.isSpecialType())
                err(e, "Cannot apply operator `%s` on types `%s` and `%s`", e.operator, t1, t2);
            return e.setInferredType(Type.BOOL_TYPE);
            default:
                return e.setInferredType(OBJECT_TYPE);
        }
    }

    @Override
    public Type analyze(Identifier id) {
        String varName = id.name;
        Type varType = currentScope.get(varName);
        if(varType!=null)
        {
            if(assign==true && !currentScope.getDeclaredSymbols().contains(varName))
                err(id, "Cannot assign to variable that is not explicitly declared in this scope: %s", varName);
            else if(assign==false && currentScope.get(varName)==null)
                err(id, "Variable not declared in scope: %s", varName);
        }
        if (varType != null && varType.isValueType()) {
            return id.setInferredType(varType);
        }

        err(id, "Not a variable: %s", varName);
        return id.setInferredType(ValueType.OBJECT_TYPE);
    }

    @Override
    public Type analyze(GlobalDecl node) 
    {
        Type ty = sym.get(node.variable.name);
        if (sym.declares(node.variable.name)==false || !isVariableType(ty)) 
        {
            err(
                node.variable, "Not a global variable: %s", node.variable.name);
            return null;
        }
        else if(currentScope.getDeclaredSymbols().contains(node.variable.name)){
            err(
                    node.variable, "Duplicate declaration of identifier in same scope: %s", node.variable.name);
        }
        else
            currentScope.put(node.variable.name, ty);
        return ty;
    }

}
