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 = new SymbolTable<>();
    /** Global symbol table. */
    private final SymbolTable<SymbolType> globals = sym;
    /** Collector for errors. */
    private Errors errors;

    /** 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) {
        sym = globalSymbols;
        errors = errors0;
    }

    /** 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) {
        errors.semError(node, message, args);
    }

    @Override
    public SymbolType analyze(VarDef varDef) {
        SymbolType lType = varDef.value.dispatch(this);
        SymbolType vType = varDef.var.dispatch(this);
        if(!vType.equals(lType))
            err(varDef, "incorrect initialization type");
        return lType;
    }

    @Override
    public SymbolType analyze(ClassDef classDef) {
        assert classDef.memberMap == null;
        classDef.memberMap = new HashMap<>();
        assert classDef.memberMap != null;

        ClassValueType result_type = new ClassValueType(classDef.name.name);
        SymbolTable<SymbolType> symLayer = new SymbolTable<>(sym);
        symLayer.put("self", result_type);
        sym = symLayer;

        for(Declaration decl : classDef.declarations) {
            String name = decl.getIdentifier().name;
            SymbolType type = decl.dispatch(this);
            classDef.memberMap.put(name, type);
        }

        sym = sym.getParent();
        return new ClassValueType(classDef.name.name);
    }

    @Override
    public SymbolType analyze(FuncDef funcDef) {
        System.out.println("debug: start funcDef: " + funcDef.getClass().getName());
        assert funcDef.symTable == null;
        sym = funcDef.symTable = new SymbolTable<SymbolType>(sym);
        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 = ValueType.annotationToValueType(param.type);
            sym.put(name, type);
            args.add(type);
        }

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

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

    @Override
    public SymbolType analyze(GlobalDecl globalDecl) {
        String name = globalDecl.variable.name;
        SymbolType T = globals.get(name);
        if(T == null)
            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.");
        //sym.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)
            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.");
        return T;
    }

    // Just a helper..
    @Override
    public SymbolType analyze(TypedVar node) {
        return ValueType.annotationToValueType(node.type);
    }

    ///////// end

    @Override
    public SymbolType analyze(Program program) {
        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.
            sym.put(name, type);
        }
        for (Stmt stmt : program.statements) {
            stmt.dispatch(this);
        }
        return null;
    }

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

    @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);
    }

    @Override
    public SymbolType analyze(AssignStmt node) {
        SymbolType right_type = node.value.dispatch(this);
        if(right_type == null)
            right_type = NONE_TYPE;
        if(node.targets.size() == 1) {
            String target = node.targets.get(0).toString();
            System.out.println("debug: assign to target " + target);
            SymbolType left_type = sym.get(target);
            if(left_type == null) {
                err(node, "Syntax error: assign to non-declared variable.");
                return OBJECT_TYPE;
            }
            if(!left_type.equals(right_type)) {
                err(node, "Syntax error: implicit type convertion not allowed.");
            }
            return right_type;
        }
        else {
            if(!right_type.isListType()) {
                err(node, "Syntax error: assign non-list type to list targets");
                return OBJECT_TYPE;
            }
            ListValueType right_ls_type = (ListValueType) right_type;
            // TODO: assign list to list
            err(node, "Not implemented error: assign list to list");
            return OBJECT_TYPE;
        }
    }

    @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 funcType = sym.get(node.function.name);
        if(funcType == null) {
            err(node, "calling undeclared function " + node.function.name);
            return OBJECT_TYPE;
        }
        if(!funcType.isFuncType()) {
            err(node, "calling non-function variable " + node.function.name);
            return OBJECT_TYPE;
        }
        FuncType fType = (FuncType)funcType;
        List<SymbolType> args_type = new ArrayList<SymbolType>();
        for(Expr arg : node.args) {
            args_type.add(arg.dispatch(this));
        }
        if(!fType.parameters.equals(args_type)) {
            err(node, "function parameter type list mismatch: " + node.function.name + ". Expected " + fType.parameters.toString() + ", got " + args_type.toString());
        }
        return node.setInferredType(fType.returnType);
    }

    @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 then_type = node.thenExpr.dispatch(this);
        SymbolType else_type = node.elseExpr.dispatch(this);
        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(!this_type.equals(elementType))
                err(node, "list elements must have the same type.");
        }
        ListValueType resultType = new ListValueType(elementType);
        return node.setInferredType(resultType);
    }

    @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;
        }
        SymbolType resultType = objType.memberMap.get(node.member.name); // currently won't work. classDef not analyzed.
        if(resultType == null) {
            err(node, "undefined identifier " + node.member.name + " as object member.");
            return NONE_TYPE;
        }
        return resultType;
    }

    @Override
    public SymbolType analyze(MethodCallExpr node) {
        FuncType funcType = (FuncType) node.method.dispatch(this);
        if(funcType == null) {
            err(node, "method is not callable.");
            return NONE_TYPE;
        }
        List<SymbolType> args_type = new ArrayList<>();
        for(Expr expr : node.args) {
            args_type.add(expr.dispatch(this));
        }
        if(!args_type.equals(funcType.parameters)) {
            err(node, "function/method parameter list mismatch");
        }
        return funcType.returnType;
    }

    @Override
    public SymbolType analyze(Identifier id) {
        String varName = id.name;
        SymbolType varType = sym.get(varName);
        System.out.println("var " + varName + "->" + varType.toString());

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

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