package chocopy.pa3;

import chocopy.common.analysis.AbstractNodeAnalyzer;
import chocopy.common.analysis.SymbolTable;
import chocopy.common.analysis.types.FuncType;
import chocopy.common.analysis.types.Type;
import chocopy.common.astnodes.*;
import chocopy.common.codegen.*;

import java.util.List;

import chocopy.common.codegen.RiscVBackend.Register;
import static chocopy.common.codegen.RiscVBackend.Register.*;

/**
 * This is where the main implementation of PA3 will live.
 *
 * <p>A large part of the functionality has already been implemented in the base class, CodeGenBase.
 * Make sure to read through that class, since you will want to use many of its fields and utility
 * methods in this class when emitting code.
 *
 * <p>Also read the PDF spec for details on what the base class does and what APIs it exposes for
 * its sub-class (this one). Of particular importance is knowing what all the SymbolInfo classes
 * contain.
 */
public class CodeGenImpl extends CodeGenBase {

    /** A code generator emitting instructions to BACKEND. */
    public CodeGenImpl(RiscVBackend backend) {
        super(backend);
    }

    /** Operation on None. */
    private final Label errorNone = new Label("error.None");
    /** Division by zero. */
    private final Label errorDiv = new Label("error.Div");
    /** Index out of bounds. */
    private final Label errorOob = new Label("error.OOB");

    /**
     * Emits the top level of the program.
     *
     * <p>This method is invoked exactly once, and is surrounded by some boilerplate code that: (1)
     * initializes the heap before the top-level begins and (2) exits after the top-level ends.
     *
     * <p>You only need to generate code for statements.
     *
     * @param statements top level statements
     */
    protected void emitTopLevel(List<Stmt> statements) {
        StmtAnalyzer stmtAnalyzer = new StmtAnalyzer(null);
        backend.emitADDI(
                SP, SP, -2 * backend.getWordSize(), "Saved FP and saved RA (unused at top level).");
        backend.emitSW(ZERO, SP, 0, "Top saved FP is 0.");
        backend.emitSW(ZERO, SP, 4, "Top saved RA is 0.");
        backend.emitADDI(FP, SP, 2 * backend.getWordSize(), "Set FP to previous SP.");

        for (Stmt stmt : statements) {
            stmt.dispatch(stmtAnalyzer);
        }
        backend.emitLI(A0, EXIT_ECALL, "Code for ecall: exit");
        backend.emitEcall(null);
    }

    /**
     * Emits the code for a function described by FUNCINFO.
     *
     * <p>This method is invoked once per function and method definition. At the code generation
     * stage, nested functions are emitted as separate functions of their own. So if function `bar`
     * is nested within function `foo`, you only emit `foo`'s code for `foo` and only emit `bar`'s
     * code for `bar`.
     */
    protected void emitUserDefinedFunction(FuncInfo funcInfo) {
        backend.emitGlobalLabel(funcInfo.getCodeLabel());
        StmtAnalyzer stmtAnalyzer = new StmtAnalyzer(funcInfo);

        for (Stmt stmt : funcInfo.getStatements()) {
            stmt.dispatch(stmtAnalyzer);
        }

        backend.emitMV(A0, ZERO, "Returning None implicitly");
        backend.emitLocalLabel(stmtAnalyzer.epilogue, "Epilogue");

        // FIXME: {... reset fp etc. ...}
        backend.emitJR(RA, "Return to caller");
    }

    /** An analyzer that encapsulates code generation for statements. */
    private class StmtAnalyzer extends AbstractNodeAnalyzer<Register> {
        /*
         * The symbol table has all the info you need to determine
         * what a given identifier 'x' in the current scope is. You can
         * use it as follows:
         *   SymbolInfo x = sym.get("x");
         *
         * A SymbolInfo can be one the following:
         * - ClassInfo: a descriptor for classes
         * - FuncInfo: a descriptor for functions/methods
         * - AttrInfo: a descriptor for attributes
         * - GlobalVarInfo: a descriptor for global variables
         * - StackVarInfo: a descriptor for variables allocated on the stack,
         *      such as locals and parameters
         *
         * Since the input program is assumed to be semantically
         * valid and well-typed at this stage, you can always assume that
         * the symbol table contains valid information. For example, in
         * an expression `foo()` you KNOW that sym.get("foo") will either be
         * a FuncInfo or ClassInfo, but not any of the other infos
         * and never null.
         *
         * The symbol table in funcInfo has already been populated in
         * the base class: CodeGenBase. You do not need to add anything to
         * the symbol table. Simply query it with an identifier name to
         * get a descriptor for a function, class, variable, etc.
         *
         * The symbol table also maps nonlocal and global vars, so you
         * only need to lookup one symbol table and it will fetch the
         * appropriate info for the var that is currently in scope.
         */

        /** Symbol table for my statements. */
        private final SymbolTable<SymbolInfo> sym;

        /** Label of code that exits from procedure. */
        protected final Label epilogue;

        /** The descriptor for the current function, or null at the top level. */
        private final FuncInfo funcInfo;
        private final String size_label;
        private int sp_off, max_sp;

        /** An analyzer for the function described by FUNCINFO0, which is null for the top level. */
        StmtAnalyzer(FuncInfo funcInfo0) {
            funcInfo = funcInfo0;
            if (funcInfo == null) {
                sym = globalSymbols;
                sp_off = max_sp = 2;
                size_label = "@..main.size";
            } else {
                sym = funcInfo.getSymbolTable();
                sp_off = max_sp = funcInfo0.getLocals().size() + 2;
                size_label = "@"+funcInfo0.getFuncName()+".size";
            }
            epilogue = generateLocalLabel();
        }

        public Register analyze(AssignStmt node) {
            return null;
        }
    
        public Register analyze(BinaryExpr node) {
            return null;
        }
    
        public Register analyze(BooleanLiteral node) {
            return null;
        }
    
        public Register analyze(CallExpr node) {
            SymbolInfo Ty = globalSymbols.get(node.function.name);
            if(Ty instanceof ClassInfo){
                //object create
                ClassInfo cls = (ClassInfo) Ty;
                /**
                    la a0, $DoublingVector$prototype         # Load pointer to prototype of: DoublingVector
                    jal alloc                                # Allocate new object in A0
                    sw a0, -12(fp)                           # Push on stack slot 3
                    sw a0, -16(fp)                           # Push argument 0 from last.
                    addi sp, fp, -16                         # Set SP to last argument.
                    lw a1, 8(a0)                             # Load address of object's dispatch table
                    lw a1, 0(a1)                             # Load address of method: DoublingVector.__init__
                    jalr a1                                  # Invoke method: DoublingVector.__init__
                    addi sp, fp, -@..main.size               # Set SP to stack frame top.
                    lw a0, -12(fp)                           # Pop stack slot 3
                 */
                backend.emitLA(A0, cls.getPrototypeLabel(), 
                    String.format("Load pointer to prototype of: %s", cls.getClassName()));
                backend.emitJAL(objectAllocLabel, "Allocate new object in A0");
                backend.emitSW(A0, FP, -sp_off*wordSize, String.format("Push on stack slot %d", sp_off));
                if(sp_off>max_sp)
                    max_sp = sp_off;
                sp_off++;
                backend.emitSW(A0, FP, -sp_off*wordSize, "Push argument 0 from last.");
                backend.emitADDI(SP, FP, sp_off, "Set SP to last argument.");
                backend.emitLW(A1, A0, getDispatchTableOffset(), "Load address of object's dispatch table");
                backend.emitLW(A1, A1, getMethodOffset(cls, "__init__"), String.format("Load address of method: %s.__init__", cls.getClassName()));
                backend.emitJALR(A1, String.format("Invoke method: %s.__init", cls.getClassName()));
                backend.emitADDI(SP, FP, "-"+size_label, "Set SP to stack frame top.");
                sp_off--;
                backend.emitLW(A0, FP, -sp_off*wordSize, String.format("Pop stack slot %d", sp_off));
            } else {
                //func call
            }

            return null;
        }

        public Register analyze(ExprStmt node) {
            return null;
        }
    
        public Register analyze(ForStmt node) {
            return null;
        }
    
        public Register analyze(Identifier node) {
            return null;
        }
    
        public Register analyze(IfExpr node) {
            return null;
        }
    
        public Register analyze(IfStmt node) {
            return null;
        }
    
        public Register analyze(IndexExpr node) {
            return null;
        }
    
        public Register analyze(IntegerLiteral node) {
            return null;
        }
    
        public Register analyze(ListExpr node) {
            return null;
        }
    
        public Register analyze(MemberExpr node) {
            ClassInfo objectClass = (ClassInfo) globalSymbols.get(node.object.getInferredType().className());
            Label label = generateLocalLabel();
            Register obj = node.object.dispatch(this);

            backend.emitBNEZ(obj, label, "Ensure not None");
            backend.emitJ(errorNone, "Go to error handler");
            backend.emitLocalLabel(label, "Not None");
            backend.emitLW(A0, obj, getAttrOffset(objectClass, node.member.name),
                 String.format("Get attribute: %s.%s",  objectClass.getClassName(), node.member.name));
            return A0;
        }
    
        public Register analyze(MethodCallExpr node) {
            Register obj = node.method.object.dispatch(this);
            int n_args = node.args.size();
            
            Label label = generateLocalLabel();
            backend.emitBNEZ(obj, label, "Ensure not None");
            backend.emitJ(errorNone, "Go to error handler");
            backend.emitLocalLabel(label, "Not None");
            if(sp_off>max_sp)
                max_sp = sp_off;
            sp_off += (n_args+1)*wordSize;
            backend.emitSW(obj, FP, (n_args - sp_off) *wordSize, String.format("Push argument %d from last.", n_args));

            for (int i = 0; i < n_args; ++i) 
                backend.emitSW(node.args.get(i).dispatch(this), FP, (n_args - i - 1 - sp_off) * wordSize, 
                    String.format("Push argument %d from last.", n_args - i - 1));
            
            backend.emitLW(A0, FP, (n_args- sp_off) * wordSize, String.format("Peek stack slot %d", sp_off - (n_args + 1)));
            ClassInfo objectClass = (ClassInfo)sym.get(((Identifier)node.method.object).name);

            backend.emitLW(A1, A0, getDispatchTableOffset(), "Load address of object's dispatch table");
            backend.emitLW(A1, A1, getMethodOffset(objectClass, node.method.member.name), 
                String.format("Load address of method: %s.%s", objectClass.getClassName(), node.method.member.name));
            backend.emitADDI(SP, FP, -sp_off * wordSize, "Set SP to last argument.");
            backend.emitJALR(A1, String.format("Invoke method: %s.%s", objectClass.getClassName(), node.method.member.name));
            backend.emitInsn(String.format("addi sp, fp, -%s", size_label), "Set SP to stack frame top.");
            sp_off -= (n_args+1)*wordSize;
            
            return A0;
        }
    
        public Register analyze(NoneLiteral node) {
            return null;
        }
    
        public Register analyze(ReturnStmt node) {
            return null;
        }
    
        public Register analyze(StringLiteral node) {
            return null;
        }
    
        public Register analyze(UnaryExpr node) {
            return null;
        }
    
        public Register analyze(WhileStmt node) {
            return null;
        }
    
        // FIXME: More, of course.

    }

    /**
     * Emits custom code in the CODE segment.
     *
     * <p>This method is called after emitting the top level and the function bodies for each
     * function.
     *
     * <p>You can use this method to emit anything you want outside of the top level or functions,
     * e.g. custom routines that you may want to call from within your code to do common tasks. This
     * is not strictly needed. You might not modify this at all and still complete the assignment.
     *
     * <p>To start you off, here is an implementation of three routines that will be commonly needed
     * from within the code you will generate for statements.
     *
     * <p>The routines are error handlers for operations on None, index out of bounds, and division
     * by zero. They never return to their caller. Just jump to one of these routines to throw an
     * error and exit the program. For example, to throw an OOB error: backend.emitJ(errorOob, "Go
     * to out-of-bounds error and abort");
     */
    protected void emitCustomCode() {
        emitErrorFunc(errorNone, "Operation on None");
        emitErrorFunc(errorDiv, "Division by zero");
        emitErrorFunc(errorOob, "Index out of bounds");
    }

    /** Emit an error routine labeled ERRLABEL that aborts with message MSG. */
    private void emitErrorFunc(Label errLabel, String msg) {
        backend.emitGlobalLabel(errLabel);
        backend.emitLI(A0, ERROR_NONE, "Exit code for: " + msg);
        backend.emitLA(A1, constants.getStrConstant(msg), "Load error message as str");
        backend.emitADDI(
                A1, A1, getAttrOffset(strClass, "__str__"), "Load address of attribute __str__");
        backend.emitJ(abortLabel, "Abort");
    }
}