package vaporize;

import cs132.vapor.ast.*;
import misc.*;

import java.util.*;

public class VaporizeVisitor extends VInstr.VisitorP<LIRDict, RuntimeException> {

    private ArrayList<String> vaporm;
    private Stack<String> spill_stack;
    private LIRDict curr;

    String[] callee_save = new String[] { "$s0", "$s1", "$s2", "$s3", "$s4", "$s5", "$s6", "$s7" };
    String[] caller_save = new String[] { "$t3", "$t4", "$t5", "$t6", "$t7", "$t8" };
    String[] arg_pass = new String[] { "$a0", "$a1", "$a2", "$a3" };
    String[] ret_pass = new String[] { "$v0" };
    String[] spillers = new String[] {"$t0", "$t1", "$t2" };


    public VaporizeVisitor(VaporProgram vp, ArrayList<String> vaporm, ArrayList<LIRDict> interval_list) {
        this.vaporm = vaporm;
        this.spill_stack = new Stack<String>();

        // is it possible for interval_list and the
        // function in vp to be out of order?
        for (int i = 0; i < vp.functions.length; ++i) {
            this.addVaporm(String.format("func %s [in %d, out %d, local %d]",
                                         vp.functions[i].ident, Math.max(vp.functions[i].params.length-4, 0),
                                         interval_list.get(i).getOutNum(),
                                         interval_list.get(i).getSpilledNum() + 14));
            for (int j = 0; j < this.callee_save.length; ++j) {
                this.addVaporm(String.format("  local[%s] = %s",
                                             j,
                                             this.callee_save[j]));
            }

            for (int j = 0; j < vp.functions[i].params.length; ++j)
                if (j < 4) { // is in a registers
                    this.addVaporm(String.format("  %s = %s",
                                                 interval_list.get(i).getInterval(vp.functions[i].params[j].toString())
                                                 .getAssignedRegister(),
                                                 arg_pass[j]));
                } else // is in 'out'
                    this.addVaporm(String.format("  %s = in[%s]",
                                                 interval_list.get(i).getInterval(vp.functions[i].params[j].toString())
                                                 .getAssignedRegister(),
                                                 j-4));

            TreeSet<Node> f = this.sortFunction(vp.functions[i]);
            MinimalLogger.info(String.format("Starting loop with function:\n %s",
                                             f.toString()));

            for (Node n : f) {
                if (n instanceof VInstr) {
                    this.curr = interval_list.get(i);
                    ((VInstr) n).accept(interval_list.get(i), this);
                    this.emptySpillStack();
                }
                else
                    this.addVaporm(((VCodeLabel) n).ident + ":");
            }
        }
    }

    public TreeSet<Node> sortFunction(VFunction f) {
        TreeSet<Node> sort = new TreeSet<Node>((v1, v2) -> {
                return Integer.compare(v1.sourcePos.line, v2.sourcePos.line);
        });

        for (VInstr s : f.body) {
            sort.add(s);
        }
        for (VCodeLabel l : f.labels) {
            sort.add(l);
        }
        return sort;
    }

    public String getRegister(String s) {
        if (this.curr.getInterval(s) == null)
            return null;
        String r = this.curr.getInterval(s).getAssignedRegister();
        if (r.contains("local")) {// SPILL
            int i = this.spill_stack.size();
            String reg = this.spillers[i];
            String str = String.format("  %s = %s",
                                       reg,
                                       r);
            MinimalLogger.info(String.format("Adding string:\n%s",
                                             str));
            this.vaporm.add(str);
            this.spill_stack.push(String.format("  %s = %s",
                                                r,
                                                reg));
            r = reg;
        }
        return r;
    }

    public void emptySpillStack() {
        for (int i = 0; i < this.spill_stack.size(); ++i)
            this.addVaporm(this.spill_stack.pop());
    }

    public ArrayList<String> getVaporm() {
        return this.vaporm;
    }

    public void addVaporm(String str) {
        MinimalLogger.info(String.format("Adding string:\n%s",
                                         str));
        this.vaporm.add(str);
        this.emptySpillStack();
    }

    public void visit(LIRDict d, VAssign n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String dest = this.getRegister(((VVarRef.Local) n.dest).ident);
        String source = (n.source instanceof VVarRef.Local) ?
            this.getRegister(n.source.toString())
            : n.source.toString();
        this.addVaporm(String.format("  %s = %s",
                                     dest, source));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(LIRDict d, VCall n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////

        for (int i = 0; i < this.caller_save.length; ++i) {
            this.addVaporm(String.format("  local[%s] = %s",
                                         i+8,
                                         this.caller_save[i]));
        }

        for (int i = 0; i < n.args.length; ++i) {
            String reg = (n.args[i] instanceof VVarRef.Local) ?
                this.getRegister(n.args[i].toString())
                : n.args[i].toString();
            if (i < 4) {// we have registers
                MinimalLogger.info(String.format("Adding argument for %s",
                                                 n.args[i].toString()));
                this.addVaporm(String.format("  %s = %s",
                                             this.arg_pass[i],
                                             reg));
            } else { // use 'out'
                this.addVaporm(String.format("  out[%s] = %s",
                                             i-4,
                                             reg));
            }
        }

        MinimalLogger.severe(String.format("n addr: %s",
                                           n.addr.toString()));
        MinimalLogger.severe(String.format("Interval: %s",
                                           this.getRegister(n.addr.toString())));
        String r = this.getRegister(n.addr.toString());
        if (r == null)
            // a label
            this.addVaporm(String.format("  call %s",
                                         n.addr.toString()));
        else
            this.addVaporm(String.format("  call %s",
                                         r));
        for (int i = 0; i < this.caller_save.length; ++i) {
            this.addVaporm(String.format("  %s = local[%s]",
                                         this.caller_save[i],
                                         i+8));
        }

        // get dest
        if (n.dest != null) {
            this.addVaporm(String.format("  %s = $v0",
                                         this.getRegister(((VVarRef.Local) n.dest).ident)
                                         ));
        }
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(LIRDict d, VBuiltIn n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String ret = "";
        if (n.dest != null) {
            ret += String.format("%s = ",
                                 this.getRegister(((VVarRef.Local) n.dest).ident)
                                 );
        }
        ret += String.format("%s(",
                             ((VBuiltIn.Op) n.op).name);

        String par;
        for (VOperand a : n.args) {
            par = (a instanceof VVarRef.Local) ?
                this.getRegister(a.toString())
                : a.toString();
            ret += String.format("%s ",
                                 par);
        }

        ret += ")";

        this.addVaporm(String.format("  %s",
                                     ret));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(LIRDict d, VMemWrite n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String dest = this.getRegister(((VMemRef.Global) n.dest).base.toString())
            ;
        String source = (n.source instanceof VVarRef.Local) ?
            this.getRegister(n.source.toString()) :
            n.source.toString();
        int byteOffset = ((VMemRef.Global) n.dest).byteOffset;
        this.addVaporm(String.format("  [%s+%d] = %s",
                                     dest, byteOffset, source));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(LIRDict d, VMemRead n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String dest = (n.dest instanceof VVarRef.Local) ?
            this.getRegister(n.dest.toString()) :
            n.dest.toString();
        String source = this.getRegister(((VMemRef.Global) n.source).base.toString())
            ;
        int byteOffset = ((VMemRef.Global) n.source).byteOffset;
        this.addVaporm(String.format("  %s = [%s+%d]",
                                     dest, source, byteOffset));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(LIRDict d, VBranch n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String ret = "  ";
        if (n.positive)
            ret += "if ";
        else
            ret += "if0 ";

        ret += this.getRegister(n.value.toString())
            + " goto :";

        ret += n.target.ident;

        this.addVaporm(String.format("%s",
                                     ret));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(LIRDict d, VGoto n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        this.addVaporm(String.format("  goto :%s",
                                     (((VAddr.Label) n.target).label).ident));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(LIRDict d, VReturn n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        // get ret
        if (n.value != null) {
            String reg = (n.value instanceof VVarRef.Local) ?
                this.getRegister(n.value.toString()) :
                n.value.toString();
            this.addVaporm(String.format("  $v0 = %s",
                                         reg));
        }
        for (int j = 0; j < this.callee_save.length; ++j) {
            this.addVaporm(String.format("  %s = local[%s]",
                                         this.callee_save[j],
                                         j));
        }
        this.addVaporm("  ret");
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

}