package condense;

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

import java.util.*;
import java.util.regex.*;

public class CondenseVisitor extends VInstr.Visitor<RuntimeException>{

    ArrayList<String> mips;

    StackHelper curr;

    public CondenseVisitor(VaporProgram vp) {
        this.mips = new ArrayList<>();

        // preamble
        this.mips.add(0, ".data");

        MinimalLogger.info("Adding data segments...");
        for (VDataSegment d : vp.dataSegments) {
            this.addMIPS(d.ident + ":");
            for (VOperand.Static s : d.values)
                this.addMIPS(String.format("  %s",
                                           s.toString().substring(1)));
        }
        this.addMIPS(".text");
        this.addMIPS("jal Main");
        this.addMIPS("li $v0 10");
        this.addMIPS("syscall");

        for (int i = 0; i < vp.functions.length; ++i) {
            this.curr = new StackHelper(vp.functions[i].stack);
            this.addMIPS(vp.functions[i].ident+":");

            // prologue
            this.addMIPS("  sw $fp -8($sp)");
            this.addMIPS("  move $fp $sp");
            this.addMIPS(String.format("  subu $sp $sp %d",
                                       this.curr.getFrameSize()));
            this.addMIPS("  sw $ra -4($fp)");

            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) {
                    ((VInstr) n).accept(this);
                }
                else
                    this.addMIPS(((VCodeLabel) n).ident + ":");
            }

        }

        this.addMIPS("_print:");
        this.addMIPS("  li $v0 1   # syscall: print integer");
        this.addMIPS("  syscall");
        this.addMIPS("  la $a0 _newline");
        this.addMIPS("  li $v0 4   # syscall: print string");
        this.addMIPS("  syscall");
        this.addMIPS("  jr $ra");

        this.addMIPS("_error:");
        this.addMIPS("  li $v0 4   # syscall: print string");
        this.addMIPS("  syscall");
        this.addMIPS("  li $v0 10  # syscall: exit");
        this.addMIPS("  syscall");

        this.addMIPS("_heapAlloc:");
        this.addMIPS("  li $v0 9   # syscall: sbrk");
        this.addMIPS("  syscall");
        this.addMIPS("  jr $ra");

        this.addMIPS(".data");
        this.addMIPS(".align 0");
        this.addMIPS("_newline: .asciiz \"\\n\"");
        this.addMIPS("_str0: .asciiz \"null pointer\\n\"");
    }

    public ArrayList<String> getMIPS() {
        return this.mips;
    }

    public void addMIPS(String str) {
        MinimalLogger.info(String.format("Adding string:\n%s",
                                         str));

        this.mips.add(str);
    }

    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 static boolean isNumeric(String strNum) {
        if (strNum == null) {
            return false;
        }
        try {
            double d = Double.parseDouble(strNum);
        } catch (NumberFormatException nfe) {
            return false;
        }
        return true;
    }

    public void visit(VAssign n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String source = n.source.toString();
        if (this.isNumeric(source))
            this.addMIPS(String.format("  li $%s %s",
                                       ((VVarRef.Register) n.dest).ident,
                                       source));
        else if (source.contains(":"))
            this.addMIPS(String.format("  la $%s %s",
                                       ((VVarRef.Register) n.dest).ident,
                                       source.substring(1)));
        else
            this.addMIPS(String.format("  move $%s %s",
                                       ((VVarRef.Register) n.dest).ident,
                                       source));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(VCall n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        if (n.addr.toString().contains(":"))
            this.addMIPS(String.format("  jal %s",
                                       n.addr.toString().substring(1)));
        else
            this.addMIPS(String.format("  jalr %s",
                                       n.addr.toString()));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(VBuiltIn n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        boolean done = false;
        String ret = "";
        String op = ((VBuiltIn.Op) n.op).name;
        String[] ts = new String[] { "$a0", "$t9"};
        // treat the arguments (place all in temp registers)

        if (!op.equals("Error")) {
            for (int i = 0; i < n.args.length; ++i) {
                if (this.isNumeric(n.args[i].toString()))
                    this.addMIPS(String.format("  li %s %s",
                                               ts[i],
                                               n.args[i].toString()));
                else
                    this.addMIPS(String.format("  move %s %s",
                                               ts[i],
                                               n.args[i].toString()));
            }
        }

        if (op.equals("Add"))
            ret += "  addu";
        else if (op.equals("Sub"))
            ret += "  subu";
        else if (op.equals("MulS"))
            ret += "  mul";
        else if (op.equals("Eq"))
            ret += "  seq";
        else if (op.equals("Lt"))
            ret += "  sltu";
        else if (op.equals("LtS"))
            ret += "  slt";
        else if (op.equals("PrintIntS")) {
            this.addMIPS("  jal _print");
            done = true;
        }
        else if (op.equals("HeapAllocZ")) {
            this.addMIPS("  jal _heapAlloc");
            if (n.dest != null)
                this.addMIPS(String.format("  move $%s $v0",
                                           ((VVarRef.Register) n.dest).ident));
            done = true;
        }
        else if (op.equals("Error")) {
            this.addMIPS("  la $a0 _str0\n  j _error");
            done = true;
        }

        if (!done) {
            if (n.dest != null)
                ret += String.format(" $%s",
                                     ((VVarRef.Register) n.dest).ident);
            for (int i = 0; i < n.args.length; ++i)
                ret += String.format(" %s",
                                     ts[i]);
            this.addMIPS(ret);
        }

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

    public void visit(VMemWrite n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String source = n.source.toString();
        if (this.isNumeric(source)) {
            this.addMIPS(String.format("  li $t9 %s",
                                       source));
            source = "$t9";
        }
        String dest;

        if (n.dest instanceof VMemRef.Stack)
            dest = this.curr.get(((VMemRef.Stack) n.dest));
        else
            dest = String.format("%d(%s)",
                                 ((VMemRef.Global) n.dest).byteOffset,
                                 ((VMemRef.Global) n.dest).base.toString());
        if (source.contains(":")) {
            this.addMIPS(String.format("  la $t9 %s",
                                       source.substring(1)));
            source = "$t9";
        }
        this.addMIPS(String.format("  sw %s %s",
                                   source,
                                   dest));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

    public void visit(VMemRead n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        String dest = n.dest.toString();
        String source;

        if (n.source instanceof VMemRef.Stack) {
            MinimalLogger.info("Loading from stack...");
            source = this.curr.get(((VMemRef.Stack) n.source));
        } else {
            MinimalLogger.info("Loading a const...");
            source = String.format("%d(%s)",
                                   ((VMemRef.Global) n.source).byteOffset,
                                   ((VMemRef.Global) n.source).base.toString());
        }
        this.addMIPS(String.format("  lw %s %s",
                                   dest,
                                   source));
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }

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

        ret += n.value.toString()
            + " ";

        ret += n.target.ident;

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

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

    public void visit(VReturn n) throws RuntimeException {
        MinimalLogger.info(String.format("->%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
        ///////////////////////////////////////////////////////////////
        this.addMIPS("  lw $ra -4($fp)");
        this.addMIPS("  lw $fp -8($fp)");
        this.addMIPS(String.format("  addu $sp $sp %d",
                                   this.curr.getFrameSize()));
        this.addMIPS("  jr $ra");
        ///////////////////////////////////////////////////////////////
        MinimalLogger.info(String.format("<-%s (%s)",
                                         n.getClass().getSimpleName(),
                                         n.sourcePos.toString()));
    }
}