package chocopy.pa2;

import chocopy.common.analysis.AbstractNodeAnalyzer;
import chocopy.common.analysis.SymbolTable;
import chocopy.common.analysis.types.*;
import chocopy.common.astnodes.*;

import java.util.*;

import static chocopy.common.analysis.types.SymbolType.*;

/** Analyzer that performs ChocoPy type checks on all nodes.  Applied after
 *  collecting declarations. */
public class TypeChecker extends AbstractNodeAnalyzer<SymbolType> {

    /** The current symbol table (changes depending on the function
     *  being analyzed). */
//    private SymbolTable<SymbolType> sym;

    /** Current symbol table.  Changes with new declarative region. */
    private SymbolTable<SymbolType> sym;
    /** Global symbol table. */
    private final SymbolTable<SymbolType> globals;
    /** Collector for errors. */
    private Errors errors;

    private boolean isBuildingSym;
    public boolean buildingSymDone;

    /** Creates a type checker using GLOBALSYMBOLS for the initial global
     *  symbol table and ERRORS0 to receive semantic errors. */
    public TypeChecker(SymbolTable<SymbolType> globalSymbols, Errors errors0, boolean firstRun) {
        sym = globalSymbols;
        globals = sym;
        errors = errors0;
        isBuildingSym = firstRun;
        buildingSymDone = true;
    }
    public SymbolTable<SymbolType> getGlobals() {
        return globals;
    }

    /** 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 void err(Node node, String message, Object... args) {
        if(!isBuildingSym)
            errors.semError(node, message, args);
    }

    ///////////////////////////// Program //////////////////////////////////////
    @Override
    public SymbolType analyze(Program program) {
        if(isBuildingSym) {
            sym.put("object", OBJECT_TYPE);
            sym.put("print", new FuncType(new ArrayList<ValueType>(Arrays.asList(OBJECT_TYPE)), NONE_TYPE));
        }

        for (Declaration decl : program.declarations) {
            String name = decl.getIdentifier().name;

            SymbolType type = decl.dispatch(this);
            System.out.println("decl type " + type.toString());

            // recolic: what's the teacher willing to do???
            if (type == null) {
                continue;
            }

            // TODO: DO NOT throw on duplicate id. generate a compiler error.
            if(isBuildingSym)
                sym.put(name, type);
        }

        if(isBuildingSym)
            // firstRun: symbolTable and memberMap is not useable now. STOP dispatching to avoid crash.
            return null;
        for (Stmt stmt : program.statements) {
            stmt.dispatch(this);
        }
        return null;
    }

    ///////////////////////////// Def&decl s //////////////////////////////////////
    @Override
    public SymbolType analyze(VarDef varDef) {
        if(isBuildingSym) {
            // firstRun: symbolTable and memberMap is not useable now. STOP dispatching to avoid crash.
            return resolveTypeAnnotation(varDef.var.type);
        }
        SymbolType lType = varDef.value.dispatch(this);
        SymbolType vType = varDef.var.dispatch(this);
        if(!(vType.equals(lType) || lType.equals(NONE_TYPE)))
            err(varDef, "incorrect initialization type");
        return vType;
    }

    @Override
    public SymbolType analyze(ClassDef classDef) {
        if(isBuildingSym) {
            if(classDef.memberMap == null)
                classDef.memberMap = new HashMap<>(); // clearing may prevent symTable from convergence forever.
        }
        assert classDef.memberMap != null;

        // My-name already pushed to symbolMap
        // TA doesn't like this: // classDef.name.dispatch(this);
        // TA doesn't like this: // classDef.superClass.dispatch(this);
        ClassValueType superType = (ClassValueType) classDef.superClass.dispatch(this);
        classDef.superClass.setInferredType(null); // make TA happy

        // Append parent to child if not exist.
        if(!classDef.superClass.name.equals("object")) {
            // This logic doesn't allow superClass use-before-decl, unless loop again and again until convergence.
            ClassValueType superClassType = (ClassValueType) sym.get(classDef.superClass.name);
            if(superClassType.memberMap != null) {
                for(String k : superClassType.memberMap.keySet()) {
                    if(!classDef.memberMap.containsKey(k))
                        classDef.memberMap.put(k, superClassType.memberMap.get(k));
                }
            }
        }

        ClassValueType result_type = new ClassValueType(classDef.name.name, superType);
        result_type.memberMap = classDef.memberMap;
        SymbolTable<SymbolType> symLayer = new SymbolTable<>(sym);
        // TODO: Maybe this is not necessary. I'm a cxx programmer.
        symLayer.put("self", result_type);
        System.out.println("debug: self_type=" + result_type.toString() + ", where symbolMap=" + result_type.memberMap);
        sym = symLayer;

        for(Declaration decl : classDef.declarations) {
            String name = decl.getIdentifier().name;
            SymbolType type = decl.dispatch(this);
            if(isBuildingSym)
                classDef.memberMap.put(name, type);
        }
        System.out.println("debug: self_type=" + result_type.toString() + ", where symbolMap=" + result_type.memberMap);

        sym = sym.getParent();
        return result_type;
    }

    @Override
    public SymbolType analyze(FuncDef funcDef) {
        if(isBuildingSym) {
            if(funcDef.symTable == null)
                funcDef.symTable = new SymbolTable<SymbolType>(sym);
            sym = funcDef.symTable;
        }
        else
            sym = funcDef.symTable;
        assert funcDef.symTable != null; // may suck if partial-compiling happens...

        // Func parameter list
        List<ValueType> args = new ArrayList<>();
        for(TypedVar param : funcDef.params) {
            String name = param.identifier.name;
            ValueType type = resolveTypeAnnotation(param.type);
            //if(typeIsUserDefinedClass(type))
            //    type = (ValueType) sym.get(((ClassValueType) type).className());
            if(isBuildingSym)
                sym.put(name, type);
            args.add(type);
        }

        for(Declaration decl : funcDef.declarations) {
            String name = decl.getIdentifier().name;
            SymbolType type = decl.dispatch(this);
            if(isBuildingSym)
                sym.put(name, type);
        }

        if(!isBuildingSym) {
            // firstRun: symbolTable and memberMap is not useable now. STOP dispatching to avoid crash.
            for(Stmt stmt : funcDef.statements) {
                stmt.dispatch(this);
            }
        }

        FuncType funcT = new FuncType(args, resolveTypeAnnotation(funcDef.returnType));
        if(isBuildingSym)
            sym.put(funcDef.name.name, funcT); // this funcdef should be add to both parentSym and localSym to support recursive func.

        // TA don't like it. OK I won't dispatch this id...
        //funcDef.name.dispatch(this); // dispatch it.

        // function scope ends.. backtracking...
        sym = sym.getParent();
        if(sym == null)
            throw new RuntimeException("logic error: sym parent is null while returning..");
        return funcT;
    }

    @Override
    public SymbolType analyze(GlobalDecl globalDecl) {
        String name = globalDecl.variable.name;
        SymbolType T = globals.get(name);
        System.out.println("DEBUG< GOL< " + isBuildingSym + " buildingSym, name=" + name + ", Type=" + T);
        if(T == null) {
            if(isBuildingSym)
                return OBJECT_TYPE; // workaround for fucking use-before-decl
            errors.semError(globalDecl, "global id '" + name + "' not found in global scope..");
        }
        if(sym == globals) // ref equal
            errors.semError(globalDecl, "global declaration '" + name + "' not allowed in global scope.");
        if(!isBuildingSym)
            // dirty! see comment in StudentAnalysis.java
            sym.overwrite_put(name, T);
        return T;
    }

    @Override
    public SymbolType analyze(NonLocalDecl nonLocalDecl) {
        String name = nonLocalDecl.variable.name;
        SymbolType T = sym.get(name); // auto-iterate through the tree.
        if(T == null) {
            if(isBuildingSym)
                return OBJECT_TYPE;
            errors.semError(nonLocalDecl, "nonlocal id '" + name + "' not found in parent scope..");
        }
        if(sym == globals)
            errors.semError(nonLocalDecl, "nonlocal declaration '" + name + "' not allowed in global scope.");
        if(!isBuildingSym)
            // dirty! see comment in StudentAnalysis.java
            sym.overwrite_put(name, T);
        return T;
    }

    ///////// end

    ///////////////////////////// Literals //////////////////////////////////////
    @Override
    public SymbolType analyze(IntegerLiteral node) {
        return node.setInferredType(INT_TYPE);
    }
    @Override
    public SymbolType analyze(StringLiteral node) {
        return node.setInferredType(STR_TYPE);
    }
    @Override
    public SymbolType analyze(NoneLiteral node) {
        return node.setInferredType(NONE_TYPE);
    }
    @Override
    public SymbolType analyze(BooleanLiteral node) {
        return node.setInferredType(BOOL_TYPE);
    }

    ///////////////////////////// Statements //////////////////////////////////////
    @Override
    public SymbolType analyze(AssignStmt node) {
        SymbolType right_type = node.value.dispatch(this);
        SymbolType left_type = null;
        if(right_type == null)
            right_type = NONE_TYPE;
        for(Expr target : node.targets) {
            System.out.println("assign: " + target.getClass().getName());
            if(target instanceof IndexExpr) {
                IndexExpr indexExpr = (IndexExpr) target;
                left_type = indexExpr.dispatch(this);
            }
            else if(target instanceof Identifier) {
                Identifier id = (Identifier) target;
                left_type = id.dispatch(this);
            }
            else if(target instanceof MemberExpr) {
                MemberExpr memberExpr = (MemberExpr) target;
                left_type = memberExpr.dispatch(this);
            }
            else {
                err(node, "Not implemented assignStmt, target=" + target.getClass().getName());
            }
            if(left_type == null) {
                err(node, "Syntax error: assign to non-declared variable");
                return NONE_TYPE;
            }
            if(!typeConvertible(right_type,  left_type)) {
                err(node, "Syntax error: implicit type convertion failed.");
            }
        }
        //return left_type;
        return null; // silly ta is not willing to use `a = (b = c) => a = something` reduction. It's ugly and buggy and complex but let him go.
    }

    @Override
    public SymbolType analyze(ForStmt node) {
        SymbolType iterType = node.iterable.dispatch(this);
        if(!(iterType.isListType() || iterType.equals(STR_TYPE))) {
            err(node, "for target is not iterable.");
        }
        SymbolType ident = node.identifier.dispatch(this);

        for(Stmt stmt : node.body) {
            stmt.dispatch(this);
        }
        return null;
    }

    @Override
    public SymbolType analyze(IfStmt node) {
        SymbolType conditionType = node.condition.dispatch(this);
        if(!conditionType.equals(BOOL_TYPE)) {
            err(node, "if statement condition should be BOOL");
        }
        for(Stmt stmt : node.thenBody) {
            stmt.dispatch(this);
        }
        for(Stmt stmt : node.elseBody) {
            stmt.dispatch(this);
        }
        return null;
    }

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

    @Override
    public SymbolType analyze(ReturnStmt node) {
        node.value.dispatch(this);
        return null;
    }

    @Override
    public SymbolType analyze(WhileStmt node) {
        SymbolType conditionType = node.condition.dispatch(this);
        if(!conditionType.equals(BOOL_TYPE)) {
            err(node, "while statement condition should be BOOL");
        }
        for(Stmt stmt : node.body) {
            stmt.dispatch(this);
        }
        return null;
    }

    ///////////////////////////// Expressions //////////////////////////////////////
    @Override
    public SymbolType analyze(BinaryExpr node) {
        String op = node.operator;
        SymbolType left_type = node.left.dispatch(this);
        SymbolType right_type = node.right.dispatch(this);
        if(left_type == null)
            left_type = NONE_TYPE;
        if(right_type == null)
            right_type = NONE_TYPE;

        switch(op) {
            case "-":
            case "*":
            case "//":
            case "%":
            case ">":
            case "<":
            case ">=":
            case "<=":
                if(!(left_type.equals(INT_TYPE) && right_type.equals(INT_TYPE))) {
                    System.out.println("debug: left=" + left_type.toString() + ", right=" + right_type.toString());
                    err(node, "Syntax Error: operand should be INT");
                }
                break;
            case "and":
            case "or":
                if(!(left_type.equals(BOOL_TYPE) && right_type.equals(BOOL_TYPE))) {
                    err(node, "Syntax Error: operand should be BOOL");
                }
                break;
            case "+":
                if(!(left_type.equals(INT_TYPE) || left_type.equals(STR_TYPE) || left_type.isListType())) {
                    err(node, "Syntax Error: operand of + should be INT or STR or LIST");
                }
                // fallthrough: do left-right match check.
            case "==":
            case "!=":
            case "is":
                if(!left_type.equals(right_type)) {
                    if(left_type.isListType() && right_type.isListType()) {
                        System.out.println("debug: left_type=" + left_type.toString() + ", right=" + right_type.toString());
                        left_type = new ListValueType(OBJECT_TYPE);
                        break; // list-type mismatch is OK. but you should return [Object]
                    }
                    err(node, "Syntax Error: binary operator operand type mismatch: left=" + left_type.toString() + ", right=" + right_type.toString());
                }
                break;
            default:
                err(node, "Syntax Error: binary operator operand type not supported.");
                return OBJECT_TYPE;
        }

        // Now set target type.
        SymbolType result_type = NONE_TYPE;
        switch(op) {
            case "-":
            case "*":
            case "//":
            case "%":
            case "and":
            case "or":
            case "+":
                result_type = left_type;
                break;
            case ">":
            case "<":
            case ">=":
            case "<=":
            case "==":
            case "!=":
            case "is":
                result_type = ValueType.BOOL_TYPE;
                break;
            default:
                throw new RuntimeException("Logic error");
        }

        return node.setInferredType(result_type);
    }

    @Override
    public SymbolType analyze(CallExpr node) {
        SymbolType calledType = node.function.dispatch(this);
        if(calledType == null) {
            err(node, "calling undeclared function " + node.function.name);
            return null;
        }

        if(calledType.isFuncType()) {
            FuncType fType = (FuncType) calledType;
            List<SymbolType> args_type = new ArrayList<SymbolType>();
            for(Expr arg : node.args) {
                args_type.add(arg.dispatch(this));
            }
            if(!matchArgList(args_type, fType.parameters)) {
                err(node, "function parameter type list mismatch: " + node.function.name + ". Expected " + fType.parameters.toString() + ", got " + args_type.toString());
            }
            return node.setInferredType(fType.returnType);
        }
        else if(calledType.equals(NONE_TYPE)) {
            err(node, "can not create instance of NONE_TYPE: " + node.function.name);
            return NONE_TYPE;
        }
        else {
            // a = int(123)
            // b = classA(fuck, you)
            node.function.setInferredType(null); // TA doesn't like this too. clear it!
            return node.setInferredType(calledType);
        }
    }

    @Override
    public SymbolType analyze(UnaryExpr node) {
        String op = node.operator;
        switch(op) {
            case "not":
                if(!node.operand.dispatch(this).equals(ValueType.BOOL_TYPE)) {
                    err(node, "operand of not should be bool.");
                }
                return node.setInferredType(ValueType.BOOL_TYPE);
            case "-":
                if(!node.operand.dispatch(this).equals(INT_TYPE)) {
                    err(node, "operand of - should be int");
                }
                return node.setInferredType(INT_TYPE);
            default:
                err(node, "operator " + op + " unknown.");
                return node.setInferredType(NONE_TYPE);
        }
    }

    @Override
    public SymbolType analyze(IfExpr node) {
        SymbolType condition_type = node.condition.dispatch(this);
        SymbolType then_type = node.thenExpr.dispatch(this);
        SymbolType else_type = node.elseExpr.dispatch(this);
        if(!condition_type.equals(BOOL_TYPE)) {
            err(node, "if_expr should have BOOL_TYPE");
        }
        if(!then_type.equals(else_type)) {
            err(node, "then_type not equals to else_type. ");
        }
        return node.setInferredType(then_type);
    }

    @Override
    public SymbolType analyze(IndexExpr node) {
        SymbolType list_type = node.list.dispatch(this);
        SymbolType index_type = node.index.dispatch(this);
        if(!index_type.equals(INT_TYPE)) {
            err(node, "index_type should be int.");
        }
        if(list_type.equals(STR_TYPE)) {
            return node.setInferredType(STR_TYPE);
        }
        else if(list_type.isListType()) {
            ListValueType t = (ListValueType) list_type;
            return node.setInferredType(t.elementType);
        }
        else {
            err(node, "type " + list_type.toString() + " is not index-able.");
            return node.setInferredType(NONE_TYPE);
        }
    }

    @Override
    public SymbolType analyze(ListExpr node) {
        SymbolType elementType = null;
        for(Expr expr : node.elements) {
            SymbolType this_type = expr.dispatch(this);
            if(elementType == null)
                elementType = this_type;
            if(!typeConvertible(this_type, elementType))  {
                //err(node, "list elements must have the same type.");
                elementType = calcCommonType(this_type, elementType);
            }
        }
        if(elementType == null)
            return node.setInferredType(EMPTY_TYPE);
        else
            return node.setInferredType(new ListValueType(elementType));
    }

    @Override
    public SymbolType analyze(MemberExpr node) {
        ClassValueType objType = (ClassValueType) node.object.dispatch(this);
        if(objType == null) {
            err(node, ". operator left operand is not object type.");
            return NONE_TYPE;
        }
        //System.out.println("manual lookup: memberMap=" + ((ClassValueType)sym.get(objType.toString())).memberMap );
        System.out.println("debug: isBuildingSym=" + isBuildingSym + ", memberMap=" + objType.memberMap + ", objtype=" + objType.toString());
        // Won't work.
        SymbolType resultType = objType.memberMap.get(node.member.name);
        if(resultType == null) {
            err(node, "undefined identifier " + node.member.name + " as object member.");
            return NONE_TYPE;
        }
        return node.setInferredType(resultType);
    }

    @Override
    public SymbolType analyze(MethodCallExpr node) {
        SymbolType methodType = node.method.dispatch(this);
        if(!( methodType instanceof FuncType)) {
            err(node, "method is not callable or not found.");
            return NONE_TYPE;
        }

        FuncType funcType = (FuncType) methodType;
        List<SymbolType> args_type = new ArrayList<>();
        if(node.method instanceof MemberExpr) {
            args_type.add(((MemberExpr)node.method).object.getInferredType());
        }
        for(Expr expr : node.args) {
            args_type.add(expr.dispatch(this));
        }

        if(!matchArgList(args_type, funcType.parameters)) {
            err(node, "function/method parameter list mismatch. Got " + args_type.toString() + ",expecting " + funcType.parameters);
        }
        return node.setInferredType(funcType.returnType);
    }
    private boolean matchArgList(List<SymbolType> got, List<ValueType> expect) {
        System.out.println("debug> match arg> " + got + " -> " + expect + ", size " + got.size() + expect.size());
        if(got.size() != expect.size())
            return false;
        for(int cter = 0; cter < got.size(); ++cter) {
            System.out.println("debug> match arg> " + got.get(cter).toString() + " -> " + expect.get(cter).toString());
            if(typeConvertible(got.get(cter), expect.get(cter)))
                continue; // ok
            else
                return false;
        }
        System.out.println("debug> SUCC match arg> " + got + " -> " + expect + ", size " + got.size() + expect.size());
        return true;
    }

    ///////////////////////////// Others //////////////////////////////////////
    @Override
    public SymbolType analyze(Identifier id) {
        String varName = id.name;
        SymbolType varType = sym.get(varName);
        if(varType == null) {
            err(id, "undefined id " + varName);
            return null;
        }
        System.out.println("var " + varName + "->" + varType.toString());

        if(varType instanceof ClassValueType) {
            // TODO: This is just a workaround. I don't know why this line is required.
            //       Should debug it sometime.
            varType = fixClassTypeProblem((ValueType) varType);
        }

        if (varType != null) {
            return id.setInferredType(varType);
        }

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

    @Override
    public SymbolType analyze(TypedVar node) {
        SymbolType result_type = resolveTypeAnnotation(node.type);
        //if(typeIsUserDefinedClass(result_type))
        //    result_type = sym.get(((ClassValueType)result_type).className());
        return result_type;
    }

    @Override
    public SymbolType analyze(ClassType node) {
        // TODO: verify it.
        throw new RuntimeException("TODO: not verified.");
        //return new ClassValueType(node.className);
    }
    @Override
    public SymbolType analyze(ListType node) {
        SymbolType elementType = node.elementType.dispatch(this);
        return new ListValueType(elementType);
    }

    ///////////////////////////// helpers ///////////////////////////////
    private boolean typeConvertible(SymbolType from, SymbolType to) {
        if(from == null || to == null)
            throw new RuntimeException("debug: SymbolType can not be null");
        if(from.equals(NONE_TYPE) || to.equals(OBJECT_TYPE))
            return true;
        if(from.equals(EMPTY_TYPE)) from = new ListValueType(NONE_TYPE);
        if(to.equals(EMPTY_TYPE)) to = new ListValueType(NONE_TYPE);
        if(!from.isListType() || !to.isListType())  {
            // Cast class to parent
            if(from instanceof ClassValueType && to instanceof ClassValueType) {
                if(_typeConvertible_calcClassInheritChainCommonNode((ClassValueType) from, (ClassValueType) to).equals(to))
                    return true; // Can cast from->to
            }
            return from.equals(to);
        }

        // Convert list to list:
        ListValueType lsFrom = (ListValueType) from, lsTo =  (ListValueType) to;
        if(lsFrom.elementType == null || lsFrom.elementType.equals(NONE_TYPE))
            return true;
        return typeConvertible(lsFrom.elementType, lsTo.elementType);
    }
    private ClassValueType _typeConvertible_calcClassInheritChainCommonNode(ClassValueType l, ClassValueType r) {
        int l_depth = 1, r_depth = 1;
        ClassValueType l_curr = l, r_curr = r;
        while(!l_curr.equals(OBJECT_TYPE)) {
            ++l_depth;
            l_curr = l_curr.parent;
        }
        while(!r_curr.equals(OBJECT_TYPE)) {
            ++r_depth;
            r_curr = r_curr.parent;
        }
        l_curr = l;
        r_curr = r;
        for(;l_depth > r_depth; --l_depth)
            l_curr = l_curr.parent;
        for(;r_depth > l_depth; --r_depth)
            r_curr = r_curr.parent;
        while(!l_curr.equals(r_curr)) { // NOTE: Use reference compare
            l_curr = l_curr.parent;
            r_curr = r_curr.parent;
        }
        if(l_curr == null) return OBJECT_TYPE;
        else return l_curr;
    }
    private SymbolType calcCommonType(SymbolType another, SymbolType preferred) {
        SymbolType l = another, r = preferred;
        if(l == null || r == null)
            throw new RuntimeException("debug: SymbolType can not be null");
        if(l.equals(NONE_TYPE))
            return r;
        if(r.equals(NONE_TYPE))
            return l;
        if(typeConvertible(l, r))
            return r;
        if(typeConvertible(r, l))
            return l;
        if(another instanceof ClassValueType && preferred instanceof ClassValueType) {
            return _typeConvertible_calcClassInheritChainCommonNode((ClassValueType) another, (ClassValueType) preferred);
        }
        return OBJECT_TYPE;
    }
    private boolean typeIsUserDefinedClass(SymbolType type) {
        if(type instanceof ClassValueType) {
            if(type.equals(INT_TYPE) || type.equals(BOOL_TYPE) || type.equals(STR_TYPE) || type.equals(NONE_TYPE) || type.equals(EMPTY_TYPE) || type.equals(OBJECT_TYPE))
                return false;
            if(type.isFuncType() || type.isListType())
                return false;
            return true;
        }
        else
            return false;
    }
    private ValueType fixClassTypeProblem(ValueType valType) {
        if(typeIsUserDefinedClass(valType)) {
            ClassValueType realType = (ClassValueType)sym.get(valType.className());
            if(realType == null) {
                System.out.println("FIXME: annotation uses an unknown class type. Unable to assign member map.");
                return valType;
            }
            return realType;
        }
        else
            return valType;
    }
    private ValueType resolveTypeAnnotation(TypeAnnotation annotation) {
        return fixClassTypeProblem(ValueType.annotationToValueType(annotation));
    }

}
