package AST.Visitor;

import java.util.List;
import java.util.Map;

import AST.*;

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


    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;
    public void visit(ClassDeclSimple n) {
        currentClass = n.i.s;
        currentMethod = null;

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

        //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");
        genbin("call", "mjcalloc");

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

        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("");

        //Vtable

        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) {
    }

    // 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.vl.size() != 0){
            genbin("subq", 8*n.vl.size() + "", "%rsp");
        }

        for (int j = 0; j < n.fl.size(); j++) {
            push(args[j]);
        }

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

        n.e.accept(this);

        for (int j = n.fl.size() - 1; j >= 0 ; j--) {
            pop(args[j]);
        }


        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) {

    }

    // Exp e;
    // StatementList sl;
    public void visit(While n) {

    }

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

    // Identifier i;
    // Exp e;
    public void visit(Assign n) {

    }

    // Identifier i;
    // Exp e1,e2;
    public void visit(ArrayAssign n) {

    }

    // Exp e1,e2;
    public void visit(And n) {

    }

    // Exp e1,e2;
    public void visit(LessThan n) {

    }

    // 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) {

    }

    // Exp e;
    public void visit(ArrayLength n) {

    }

    // Exp e;
    // Identifier i;
    // ExpList el;
    public void visit(Call n) {
        n.e.accept(this);
        if(!(n.e instanceof This)) {
            genbin("movq", "%rax", "%rdi");
        }


        genbin("movq", "%rdi", "%rbx");
        for (int i = 0; i < n.el.size(); i++) {
            n.el.get(i).accept(this);

            if (n.el.get(i) instanceof This) {
                genbin("movq", "%rdi", args[i]);
            } else {
                genbin("movq", "%rax", args[i]);
            }

        }

        genbin("call", n.e.expType + "$" + n.i.s);
    }

    // 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) {

    }

    public void visit(False n) {

    }

    // String s;
    public void visit(IdentifierExp n) {
        List<ParameterTuple> methodArgs = info.get(currentClass).methods.get(currentMethod).argsOrder;
        for (int i = 0; i < methodArgs.size(); i++) {
            if (n.s.equals(methodArgs.get(i).name)) {
                genbin("movq", "-" + ((i + 1) * 8) + "(%rbp)", "%rax");
                return;
            } else {
                //save for later when we include field variables/local declarations as well

            }
        }
    }

    public void visit(This n) {
    }

    // Exp e;
    public void visit(NewArray n) {

    }

    // Identifier i;
    public void visit(NewObject n) {
        genbin("call", n.i.s + "$" + n.i.s);
//        genbin("movq", "$" + (info.get(n.i.s).fields.size() * 8), "%rax");
//        genbin("call", "mjcalloc");
//        genbin("leaq", n.i.s + "$$", "%rdx");
//        genbin("movq", "%rdx", "0(%rax)");
//        push("%rax");
//        genbin("call", n.i.s + "$" + n.i.s);
//        pop("%rax");
    }

    // Exp e;
    public void visit(Not n) {

    }

    // 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 genLabel(String L){
        gen(L + ":");
    }
}