package AST.Visitor;

import java.util.*;

import AST.*;

public class CodeGenVisitor implements Visitor {
    public int memory;
    public Map<String, ClassTable> info;
    public String currentClass;
    public String currentMethod;
    public String[] args = {"%rdi", "%rsi", "%rdx", "%rcx", "%r8", "%r9"};

    public Map<String, Integer> labels = new HashMap<>();
    public boolean aligned = false;


    public CodeGenVisitor(){

    }

    public CodeGenVisitor(Map<String, ClassTable> info){
        this.info = info;
        this.memory = 0;
    }


    // MainClass m;
    // ClassDeclList cl;
    public void visit(Program n) {
        gen("\t.text");
        gen("\t.globl asm_main");
        n.m.accept(this);

        for(int i = 0; i < n.cl.size(); i++){
            n.cl.get(i).accept(this);
        }
    }

    // Identifier i1,i2;
    // Statement s;
    public void visit(MainClass n) {
        currentClass = n.i1.s;
        currentMethod = null;

        genLabel("asm_main");

        //Prologue
        //Allocate stack frame
        push("%rbp");
        genbin("movq", "%rsp","%rbp");

        n.s.accept(this);

        //Epilogue
        genbin("movq","%rbp", "%rsp");
        pop("%rbp");
        gen("\tret");

        System.out.println();
        gen("\t.data");
        gen(n.i1.s + "$$:\t.quad\t0");

    }

    // Identifier i;
    // VarDeclList vl;
    // MethodDeclList ml;f
    public void visit(ClassDeclSimple n) {
        currentClass = n.i.s;
        currentMethod = null;

        //Constructor
        genLabel(n.i.s + "$" + n.i.s);

        push("%rbp");
        genbin("movq", "%rsp","%rbp");

        push("%rdi");

        genbin("movq", "$" + (8 + (info.get(n.i.s).fields.size() * 8)), "%rdi");

        align("%rax");
        genbin("call", "mjcalloc");
        dealign("%r14");

        pop("%rdi");
        genbin("leaq", n.i.s + "$$", "%rdx");
        genbin("movq", "%rdx", "0(%rax)");
        for (int i = 0; i < n.vl.size(); i++) {
            info.get(currentClass).fieldOffset.put(n.vl.get(i).i.s, (i + 1) * (-8));
        }

        for(int i = 8; i < (8 + info.get(n.i.s).fields.size() * 8); i += 8){
            genbin("movq", "$0", i + "(%rax)");
        }

        genbin("movq", "%rbp", "%rsp");
        pop("%rbp");
        gen("\tret");
        gen("");

        for(int i = 0; i < n.ml.size(); i++){
            currentMethod = n.ml.get(i).i.s;
            n.ml.get(i).accept(this);
        }

        gen("\t.data");
        gen(n.i.s + "$$:\t.quad\t0");
        gen( "\t\t.quad\t" + currentClass + "$" + currentClass);
        for(int i = 0; i < n.ml.size(); i++){
            gen( "\t\t.quad\t" + currentClass + "$" + n.ml.get(i).i.s);
        }
    }

    // Identifier i;
    // Identifier j;
    // VarDeclList vl;
    // MethodDeclList ml;
    public void visit(ClassDeclExtends n) {
        currentClass = n.i.s;
        currentMethod = null;
        genLabel(n.i.s + "$" + n.i.s);

        push("%rbp");
        genbin("movq", "%rsp","%rbp");

        push("%rdi");

        genbin("movq", "$" + (8 + (info.get(n.i.s).totalFields * 8)), "%rdi");

        align("%rax");
        genbin("call", "mjcalloc");
        dealign("%r14");


        pop("%rdi");
        genbin("leaq", n.i.s + "$$", "%rdx");
        genbin("movq", "%rdx", "0(%rax)");

        int fieldsBefore = info.get(info.get(currentClass).parent).totalFields;
        for (int i = 0; i < n.vl.size(); i++) {
            info.get(currentClass).fieldOffset.put(n.vl.get(i).i.s, (i + 1 + fieldsBefore) * (-8));
        }

        for(int i = 8; i < (8 + info.get(n.i.s).totalFields * 8); i += 8){
            genbin("movq", "$0", i + "(%rax)");
        }

        genbin("movq", "%rbp", "%rsp");
        pop("%rbp");
        gen("\tret");
        gen("");

        for(int i = 0; i < n.ml.size(); i++){
            currentMethod = n.ml.get(i).i.s;
            n.ml.get(i).accept(this);
        }

        gen("\t.data");
        gen(n.i.s + "$$:\t.quad\t0");
        gen( "\t\t.quad\t" + currentClass + "$" + currentClass);

        Map<String, MethodTable> allMethods = info.get(currentClass).allMethods;
        for (String s : allMethods.keySet()) {
            gen("\t\t.quad\t" + allMethods.get(s).classname + "$" + allMethods.get(s).name);
        }
    }

    // Type t;
    // Identifier i;
    public void visit(VarDecl n) {
    }

    // Type t;
    // Identifier i;
    // FormalList fl;
    // VarDeclList vl;
    // StatementList sl;
    // Exp e;
    public void visit(MethodDecl n) {
        currentMethod = n.i.s;
        genLabel(currentClass + "$" + n.i.s);

        push("%rbp");
        genbin("movq", "%rsp","%rbp");

        if (n.fl.size() == 0) {
            push("%rdi");
        } else {
            for (int i = 0; i <= n.fl.size(); i++) {
                if (i == 0) {
                    push("%rdi");
                } else {
                    push(args[i]);
                }
            }
        }
        for (int i = 0; i < n.vl.size(); i++) {
            push("$0");
        }

        for(int i = 0; i < n.sl.size(); i++){
            n.sl.get(i).accept(this);
        }

        n.e.accept(this);
        for (int i = n.vl.size() - 1; i >= 0; i--) {
            pop("%rsi");
        }
        if (n.fl.size() == 0) {
            pop("%rdi");
        } else {
            for (int i = n.fl.size(); i >= 0; i--) {
                if (i == 0) {
                    pop("%rdi");
                }
                pop(args[i]);
            }
        }

        genbin("movq","%rbp", "%rsp");
        pop("%rbp");
        gen("\tret");
    }

    // Type t;
    // Identifier i;
    public void visit(Formal n) {
    }

    public void visit(IntArrayType n) {
    }

    public void visit(BooleanType n) {
    }

    public void visit(IntegerType n) {
    }

    // String s;
    public void visit(IdentifierType n) {

    }

    // StatementList sl;
    public void visit(Block n) {
        for(int i = 0; i < n.sl.size(); i++){
            n.sl.get(i).accept(this);
        }
    }

    // Exp e;
    // StatementList s1,s2;
    public void visit(If n) {
        String _else_ = "else" + labelnum("else");
        String _done_ = "done" + labelnum("done");

        n.e.accept(this);
        genbin("cmpq", "$0", "%rax");
        genbin("je", _else_);
        n.s1.get(0).accept(this);
        genbin("jmp", _done_);
        genLabel(_else_);
        n.s2.get(0).accept(this);
        genLabel(_done_);
    }

    // Exp e;
    // StatementList sl;
    public void visit(While n) {
        String _test_ = "test" + labelnum("test");
        String _done_ = "done" + labelnum("done");

        genLabel(_test_);

        n.e.accept(this);
        genbin("cmpq", "$0", "%rax");
        genbin("je", _done_);
        n.sl.get(0).accept(this);
        genbin("jmp", _test_);
        genLabel(_done_);
    }

    // Exp e;
    public void visit(Print n) {
        n.e.accept(this);
        push("%rdi");
        genbin("movq", "%rax", "%rdi");

        align("%rax");
        genbin("call", "put");
        dealign("%r14");
        pop("%rdi");
    }

    // Identifier i;
    // Exp e;
    public void visit(Assign n) {
        n.e.accept(this);
        genbin("movq", "%rax", getParamLocalOffset(n.i.s));
    }

    // Identifier i;
    // Exp e1,e2;
    public void visit(ArrayAssign n) {
        n.e1.accept(this); //index
        push("%rax");

        n.e2.accept(this); //value
        pop("%rdx");

        //Array pointer in %rax

        genbin("movq", getParamLocalOffset(n.i.s), "%rcx");
        genbin("addq", "$1", "%rdx");

        genbin("movq", "%rax", "(%rcx, %rdx, 8)");

    }

    // Exp e1,e2;
    public void visit(And n) {
        String _and_ = "and" + labelnum("and");
        String _done_ = "done" + labelnum("done");

        n.e1.accept(this);
        genbin("cmpq", "$0", "%rax");
        genbin("je", _and_);

        n.e2.accept(this);
        genbin("cmpq", "$0", "%rax");
        genbin("je", _and_);

        genbin("movq", "$1", "%rax");
        genbin("jmp", _done_);

        genLabel(_and_);
        genbin("movq", "$0", "%rax");

        genLabel(_done_);

    }

    // Exp e1,e2;
    public void visit(LessThan n) {
        String _lt_ = "lt" + labelnum("lt");
        String _done_ = "done" + labelnum("done");

        n.e1.accept(this);
        push("%rax");
        n.e2.accept(this);
        pop("%rdx");
        genbin("cmpq", "%rax", "%rdx");
        genbin("jl", _lt_);
        genbin("movq", "$0", "%rax");
        genbin("jmp", _done_);
        genLabel(_lt_);
        genbin("movq", "$1", "%rax");
        genbin("jmp", _done_);
        genLabel(_done_);
    }

    // Exp e1,e2;
    public void visit(Plus n) {
        n.e1.accept(this);
        push("%rax");
        n.e2.accept(this);
        pop("%rdx");
        genbin("addq", "%rdx", "%rax");
    }

    // Exp e1,e2;
    public void visit(Minus n) {
        n.e2.accept(this);
        push("%rax");
        n.e1.accept(this);
        pop("%rdx");

        genbin("subq", "%rdx", "%rax");

    }

    // Exp e1,e2;
    public void visit(Times n) {
        n.e1.accept(this);
        push("%rax");
        n.e2.accept(this);
        pop("%rdx");
        genbin("imulq", "%rdx", "%rax");
    }

    // Exp e1,e2;
    public void visit(ArrayLookup n) {
        n.e1.accept(this);
        push("%rax");
        n.e2.accept(this);
        pop("%rdx");

        genbin("movq", "8(%rdx, %rax, 8)", "%rax");
    }

    // Exp e;
    public void visit(ArrayLength n) {
        n.e.accept(this);
        genbin("movq", "0(%rax)", "%rax");
    }

    // Exp e;
    // Identifier i;
    // ExpList el;
    public void visit(Call n) {

        push("%rdi");

        genbin("movq", "%rdi", "%rbx");
        n.e.accept(this);

        if(!(n.e instanceof This)) {
            genbin("movq", "%rax", "%rdi");
        }

        for (int i = 1; i <= n.el.size(); i++) {
            n.el.get(i - 1).accept(this);
            if (n.el.get(i - 1) instanceof This) {
                genbin("movq", "%rdi", args[i]);
            } else {
                genbin("movq", "%rax", args[i]);
            }
        }

        if((memory + 8 * n.el.size()) % 16 != 0){
            align("%rax");
        }

        genbin("call", info.get(n.e.expType).allMethods.get(n.i.s).classname + "$" + n.i.s);


        dealign("%r14");

        pop("%rdi");
    }

    // int i;
    public void visit(IntegerLiteral n) {
        if(n.i == 0){
            genbin("xorq", "%rax", "%rax");
        }
        else{
            genbin("movq", "$" + n.i, "%rax");
        }
    }

    public void visit(True n) {
        genbin("movq", "$1", "%rax");

    }

    public void visit(False n) {
        genbin("movq", "$0", "%rax");
    }

    // String s;
    public void visit(IdentifierExp n) {
        genbin("movq", getParamLocalOffset(n.s), "%rax");
    }

    public void visit(This n) {
        genbin("movq", "%rdi", "%rax");
    }

    // Exp e;
    public void visit(NewArray n) {
        push("%rdi");
        n.e.accept(this);

        //Allocation 8 bytes for length and 8*length bytes
<<<<<<< HEAD

        push("%rdx");
        push("%rax");

        genbin("addq", "$1", "%rdi");
        genbin("imulq", "$8", "%rdi");

        genbin("movq", "%rax", "%rdi");

        align("%rax");
        genbin("call", "mjcalloc");
        dealign("%r14");
=======

        push("%rdx");
        push("%rax");

        genbin("addq", "$1", "%rdi");
        genbin("imulq", "$8", "%rdi");

        genbin("movq", "%rax", "%rdi");
        genbin("call", "mjcalloc");
>>>>>>> 448d8422558ba182b5dfbfea370272c71a5edc2f

        pop("%rdx");

        genbin("movq", "%rdx", "0(%rax)");

        genbin("addq", "$1", "%rdx");
        genbin("imulq", "$8", "%rdx");

        String _arr_ = "arr" + labelnum("arr");
        String _done_ = "done" + labelnum("done");

        push("%rcx");


        //%rax = pointer to array
        //%rdx = array length
        //%rcx = current index
        genbin("movq", "$8", "%rcx");

        //i < 8 + arr.length * 8
        genLabel(_arr_);
        genbin("cmpq", "%rcx", "%rdx");

        genbin("je", _done_);
        genbin("movq", "$0",  "(%rcx, %rax)");
        genbin("addq", "$8", "%rcx");
        genbin("jmp", _arr_);

        genLabel(_done_);

        pop("%rcx");
        pop("%rdx");
        pop("%rdi");
    }

    // Identifier i;
    public void visit(NewObject n) {
        align("%rax");
        genbin("call", n.i.s + "$" + n.i.s);
        dealign("%r14");
    }

    // Exp e;
    public void visit(Not n) {
        n.e.accept(this);
        push("%rdx");
        push("%rax");

        genbin("movq", "$1", "%rax");
        pop("%rdx");

        genbin("subq", "%rdx", "%rax");
        pop("%rdx");
    }

    // String s;
    public void visit(Identifier n) {

    }

    // Exp e;
    public void visit(Display n) {
    }


    public void gen(String s){
        System.out.println(s);
    }

    public void genbin(String op, String src, String dst){
        gen("\t" + op + "\t" + src + "," + dst);
    }

    public void genbin(String op, String r){
        gen("\t" + op + "\t" + r);
    }

    public void push(String s){
        memory += 8;
        genbin("pushq", s);
    }

    public void pop(String s){
        memory -= 8;
        genbin("popq", s);
    }

    public void align(String s){
        if(memory % 16 != 0) {
            push(s);
            aligned = true;
        }
    }


    public void dealign(String s){
        if(aligned) {
            pop(s);
            aligned = false;
        }
    }

    public void genLabel(String L){
        gen(L + ":");
    }

    public String getParamLocalOffset(String s) {
        //parameters
        List<ParameterTuple> methodArgs = info.get(currentClass).methods.get(currentMethod).argsOrder;
        for (int i = 0; i < methodArgs.size(); i++) {
            if (s.equals(methodArgs.get(i).name)) {
                return "-" + ((i + 2) * 8) + "(%rbp)";
            }
        }
        //local declarations
        List<ParameterTuple> varsOrder = info.get(currentClass).methods.get(currentMethod).varsOrder;
        int methodOffset = 8 * (methodArgs.size() + 1);
        for (int i = 0; i < varsOrder.size(); i++) {
            if (s.equals(varsOrder.get(i).name)) {
                return "-" + (methodOffset + (i + 1) * 8) + "(%rbp)";
            }
        }
        //class
        genbin("movq", "-8(%rbp)", "%r15");

        if (info.get(currentClass).fieldOffset.containsKey(s)) {
            return info.get(currentClass).fieldOffset.get(s) + "(%r15)";
        } else {
            Integer offset = null;
            String tempClass = currentClass;
            while (offset == null) {
                tempClass = info.get(currentClass).parent;
                offset = info.get(tempClass).fieldOffset.get(s);
            }
            return offset + "(%r15)";
        }
    }

    public int labelnum(String s){
        if (!labels.containsKey(s)) {
            labels.put(s, 1);
        } else {
            labels.put(s, labels.get(s) + 1);
        }
        return labels.get(s);
    }
}