public class CodegenImpl implements Codegen {
	private int labelCounter = 0;
	
	public CodegenImpl() {}
	
	public String codegen(Program p) throws CodegenException {
		StringBuilder builder = new StringBuilder();
		Declaration main = null;

		for (Declaration d : p.decls) {
			if (main == null) {
				main = d;
			}
			
			builder.append(compileDeclaration(d));
		}
		
		return "call " + String.format("%s_%d\n", main.id, main.numOfArgs)
		+ "li a7,10\n"
		+ "ecall\n"
		+ builder.toString();
	}

	private String generateLabel() {
		return String.format("branch_%d", labelCounter++);
	}

	public String compileDeclaration(Declaration d) throws CodegenException {
		return String.format("%s_%d:\n", d.id, d.numOfArgs)
		+ compileExp(d.body)
		+ "ret\n";
	}

	public String compileExp(Exp e) throws CodegenException {
		if (e instanceof IntLiteral) {
			return String.format("li a0,%d\n", ((IntLiteral) e).n);
		} else if (e instanceof If) {
			If eIf = (If) e;

			String thenLabel = generateLabel();
			String elseLabel = generateLabel();
			String finalizeLabel = generateLabel();
			
			return compileFullComp(eIf.l, eIf.r, eIf.comp, thenLabel, elseLabel)
			+ thenLabel + ": " + compileExp(eIf.thenBody)
			+ "j " + finalizeLabel + "\n"
			+ elseLabel + ": " + compileExp(eIf.elseBody)
			+ finalizeLabel + ": " + "lw s1,(sp)\n"
			+ "addi sp,sp,4\n";
		} else if (e instanceof Binexp) {
			Binexp b = (Binexp) e;

			return compileExp(b.l)
			+ "sw s1,-4(sp)\n"
			+ "addi sp,sp,-4\n"
			+ "mv s1,a0\n"
			+ compileExp(b.r)
			+ "mv t1,a0\n"
			+ compileExpBinop(b.binop)
			+ "lw s1,(sp)\n"
			+ "addi sp,sp,4\n";
		} else if (e instanceof Variable) {
			int offset = ((Variable) e).x - 1;
			
			return String.format("lw a0,%d(s0)\n", offset * 4);
		} else if (e instanceof Invoke) {
			return compileCall((Invoke) e, false);
		} else if (e instanceof Seq) {
			String result = "";
			Seq statements = (Seq) e;

			if (statements.r instanceof Invoke) {
				// tail call
				result = compileCall((Invoke) statements.r, true);
			} else {
				result = compileExp(statements.r);
			}

			return compileExp(statements.l) + result;
		} else if (e instanceof While) {
			While eWhile = (While) e;

			String startLabel = generateLabel();
			String bodyLabel = generateLabel();
			String exitLabel = generateLabel();
			
			return "sw s2,-4(sp)\n"
			+ "sw s3,-8(sp)\n"
			+ "addi sp,sp,-8\n"
			+ "la s2," + startLabel + "\n"
			+ "la s3," + exitLabel + "\n"
			+ startLabel + ": " + compileFullComp(eWhile.l, eWhile.r, eWhile.comp, bodyLabel, exitLabel)
			+ bodyLabel + ": " + compileExp(eWhile.body)
			+ "lw s1,(sp)\n"
			+ "addi sp,sp,4\n"
			+ "j " + startLabel + "\n"
			+ exitLabel + ": lw s1,(sp)\n"
			+ "lw s3,4(sp)\n"
			+ "lw s2,8(sp)\n"
			+ "addi sp,sp,12\n";
		} else if (e instanceof RepeatUntil) {
			RepeatUntil eUntil = (RepeatUntil) e;

			String compLabel = generateLabel();
			String bodyLabel = generateLabel();
			String exitLabel = generateLabel();
			String contLabel = generateLabel();
			String breakLabel = generateLabel();

			return "sw s2,-8(sp)\n"
			+ "sw s3,-4(sp)\n"
			+ "addi sp,sp,-8\n"
			+ "la s2," + contLabel + "\n"
			+ "la s3," + breakLabel + "\n"
			+ bodyLabel + ": " + compileExp(eUntil.body)
			+ contLabel + ": sw s4,-4(sp)\n"
			+ "addi sp,sp,-4\n"
			+ "mv s4,a0\n"
			+ compileFullComp(eUntil.l, eUntil.r, eUntil.comp, exitLabel, compLabel)
			+ compLabel + ": lw s1,(sp)\n"
			+ "lw s4,4(sp)\n"
			+ "addi sp,sp,8\n"
			+ "j " + bodyLabel + "\n"
			+ breakLabel + ": sw s1,-8(sp)\n"
			+ "sw s4,-4(sp)\n"
			+ "addi sp,sp,-8\n"
			+ "mv s4,a0\n"
			+ exitLabel + ": lw s1,(sp)\n"
			+ "mv a0,s4\n" // body result to accumulator
			+ "lw s4,4(sp)\n" // restore previous s4
			+ "lw s2,8(sp)\n"
			+ "lw s3,12(sp)\n"
			+ "addi sp,sp,16\n";
		} else if (e instanceof Assign) {
			Assign assign = (Assign) e;

			return compileExp(assign.e)
			+ String.format("sw a0,%d(s0)\n", (assign.x - 1) * 4);
		} else if (e instanceof Continue) {
			return "jr s2\n";
		} else if (e instanceof Break) {
			return "jr s3\n";
		} else if (e instanceof Skip) {
			return "nop\n";
		}

		throw new CodegenException("Unsupported expression!");
	}

	public String compileCall(Invoke call, boolean tailCall) throws CodegenException {
		int argLen = call.args.size();
		String methodLabel = String.format("%s_%d", call.name, argLen);
		StringBuilder builder = new StringBuilder();

		// Store return address
		builder.append("sw ra,-4(sp)\n");
		builder.append("mv t0,sp\n");
		// shuffle stack pointer to start of arguments
		builder.append(String.format("addi sp,sp,-%d\n", (argLen + 1) * 4));
		builder.append("sw t0,-8(sp)\n"); // store old stack pointer (top of frame)

		for (int i = argLen - 1; i >= 0; i--) {
			Exp argExp = call.args.get(i);

			builder.append(compileExp(argExp));
			builder.append(
				String.format("sw a0,%d(sp)\n", i*4)
			);
		}

		// push old frame pointer
		builder.append("sw s0,-4(sp)\n");
		builder.append("addi sp,sp,-8\n");

		builder.append(
			// set frame pointer to start of arguments
			"addi s0,sp,8\n"
			+ "call " + methodLabel + "\n"
			// restore stack to frame pointer - 8 (bottom of frame)
			// this ensures any stack growth inside the method is undone
			+ "addi sp,s0,-8\n"
			// restore old frame pointer & stack pointer
			+ "lw s0,4(sp)\n"
			+ "lw sp,(sp)\n"
			// restore return address of caller
			+ "lw ra,-4(sp)\n"
		);

		return builder.toString();
	}

	public String compileFullComp(Exp left, Exp right, Comp c, String successLabel, String failLabel) throws CodegenException {
		return compileExp(left)
			+ "sw s1,-4(sp)\n"
			+ "addi sp,sp,-4\n"
			+ "mv s1,a0\n"
			+ compileExp(right)
			+ "mv t1,a0\n"
			+ compileComp(c, successLabel, failLabel);
	}

	public String compileComp(Comp c, String thenLabel, String elseLabel) throws CodegenException {
		String instruction = "";
		
		if (c instanceof Equals) {
			instruction = "beq";
		} else if (c instanceof Less) {
			instruction = "blt";
		} else if (c instanceof LessEq) {
			instruction = "ble";
		} else if (c instanceof Greater) {
			instruction = "bgt";
		} else if (c instanceof GreaterEq) {
			instruction = "bge";
		} else {
			throw new CodegenException("Unsupported comparison!");
		}

		return String.format("%s s1,t1,%s\nb %s\n", instruction, thenLabel, elseLabel);
	}

	public String compileExpBinop(Binop bin) throws CodegenException {
		if (bin instanceof Plus) {
			return "add a0,s1,t1\n";
		} else if (bin instanceof Minus) {
			return "sub a0,s1,t1\n";
		} else if (bin instanceof Times) {
			return "mul a0,s1,t1\n";
		} else if (bin instanceof Div) {
			return "div a0,s1,t1\n";
		}

		throw new CodegenException("Unsupported bin-op! " + bin);
	}
}