package pp.s1184725.boppi;

import java.util.*;
import java.util.logging.Logger;
import java.util.stream.Collectors;

import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.tree.*;
import org.apache.commons.lang3.StringUtils;

import pp.iloc.eval.Machine;
import pp.iloc.model.*;
import pp.s1184725.boppi.antlr.BoppiBaseVisitor;
import pp.s1184725.boppi.antlr.BoppiLexer;
import pp.s1184725.boppi.antlr.BoppiParser.*;

/**
 * The generator stage of the Boppi toolchain.
 */
public class BoppiGenerator extends BoppiBaseVisitor<Void> {
	private static final int ARBASESIZE = 16;
	private static final Operand ZERO = new Num(0);
	private static final Operand OFFSET_ARP = new Num(-4), OFFSET_RETURN_ADDR = new Num(-8),
			OFFSET_RETURN_VAL = new Num(-12), OFFSET_AL = new Num(-16);

	/**
	 * Recommended number of registers to use before throwing a warning.
	 */
	public static int RECOMMENDED_REGISTER_COUNT = 10;
	private Program prog;
	private Annotations an;
	private int regNum;
	private final Reg tempArp = new Reg("ART"), arp = Machine.ARP_REG;
	private List<Reg> regFree, regInUse;
	private Logger logger;

	/**
	 * Maps operator tokens to ILOC opcodes.
	 */
	private static Map<Integer, OpCode> ops = new HashMap<Integer, OpCode>() {
		private static final long serialVersionUID = 8979722313842633807L;

		{
			put(BoppiLexer.AND, OpCode.and);
			put(BoppiLexer.DIVIDE, OpCode.div);
			put(BoppiLexer.EQ, OpCode.cmp_EQ);
			put(BoppiLexer.GT, OpCode.cmp_GT);
			put(BoppiLexer.GTE, OpCode.cmp_GE);
			put(BoppiLexer.LEQ, OpCode.cmp_LE);
			put(BoppiLexer.LT, OpCode.cmp_LT);
			put(BoppiLexer.MINUS, OpCode.sub);
			put(BoppiLexer.MULTIPLY, OpCode.mult);
			put(BoppiLexer.NEQ, OpCode.cmp_NE);
			put(BoppiLexer.PLUS, OpCode.add);
			put(BoppiLexer.OR, OpCode.or);
		}
	};

	/**
	 * Creates a new generator with the given variable annotations and attaches
	 * the given logger.
	 * 
	 * @param ast
	 *            the program's AST
	 * @param annotations
	 *            the annotation object to use during generation
	 * @param logger
	 *            the logger to report errors to
	 * @return an ILOC program generated from the AST
	 */
	public static Program generateProgram(ParseTree ast, Annotations annotations, Logger logger) {
		BoppiGenerator generator = new BoppiGenerator(annotations, logger);
		generator.visit(ast);
		return generator.prog;
	}

	protected BoppiGenerator(Annotations annotations, Logger logger) {
		an = annotations;
		prog = new Program();
		regNum = 0;
		regFree = new ArrayList<>();
		regInUse = new ArrayList<>();
		this.logger = logger;
	}

	private Reg makeReg(ParserRuleContext ctx) {
		Reg reg;
		if (regFree.isEmpty()) {
			reg = new Reg("__" + (regInUse.size() + 1));
			regInUse.add(reg);

			if (regInUse.size() == 11)
				logger.warning(
						String.format("Using more than 10 registers at %d:%d. Consider rebalancing your expressions.",
								ctx.start.getLine(), ctx.start.getCharPositionInLine()));
		} else {
			if (regFree.remove(an.registers.get(ctx)))
				reg = an.registers.get(ctx);
			else
				reg = regFree.remove(0);
			regInUse.add(reg);
		}
		an.registers.put(ctx, reg);
		return reg;
	}

	private void freeReg(ParserRuleContext ctx) {
		Reg reg = an.registers.get(ctx);
		if (reg == null)
			return;
		if (!regInUse.contains(reg)) {
			logger.severe(getError(ctx, "INTERNAL: cannot free register %s since it is not in use.", reg.getName()));
		} else {
			regInUse.remove(reg);
			regFree.add(reg);
		}
	}

	/**
	 * Retrieves a register for a given tree node.
	 * 
	 * @param parseTree
	 *            the tree node
	 * @return a register for the tree node
	 */
	private Reg getReg(ParseTree parseTree) {
		Reg reg = an.registers.get(parseTree);
		if (reg == null) {
			if (Math.max(3, 4) > 2)
				throw new RuntimeException("???" + parseTree.getText());
			reg = new Reg("r_" + (regNum++));
			an.registers.put(parseTree, reg);
		}
		return reg;
	}

	/**
	 * Attaches the register used in the {@code source} tree node to the
	 * {@link destination} node.
	 * 
	 * @param source
	 *            the node to retrieve the register from
	 * @param destination
	 *            the node to copy the register to
	 * @return the register that was copied
	 */
	private Reg copyReg(ParseTree source, ParseTree destination) {
		Reg reg = an.registers.get(source);
		an.registers.put(destination, reg);
		return reg;
	}

	private Reg getVar(ParserRuleContext node) {
		Variable<Type> var = an.variables.get(node);

		if (var.getDepth() <= an.function.get(node).getDepth()) {
			emit("travelling ALs", OpCode.i2i, arp, tempArp);
			for (int i = var.getDepth(); i <= an.function.get(node).getDepth(); i++)
				emit(StringUtils.repeat(' ', i) + "\\", OpCode.loadAI, tempArp, OFFSET_AL, tempArp);

			return tempArp;
		} else
			return arp;
	}

	/**
	 * Constructs an operation from the parameters and adds it to the program
	 * under construction.
	 */
	private Op emit(String comment, Label label, OpCode opCode, Operand... args) {
		Op result = new Op(label, opCode, args);
		result.setComment(String.format("%-40s | %20s : %s", comment,
				regInUse.stream().map((reg) -> reg.getName().substring(2)).collect(Collectors.joining(" ")),
				regFree.stream().map((reg) -> reg.getName().substring(2)).collect(Collectors.joining(" "))));
		this.prog.addInstr(result);
		return result;
	}

	/**
	 * Constructs an operation from the parameters and adds it to the program
	 * under construction.
	 */
	private Op emit(String comment, OpCode opCode, Operand... args) {
		return emit(comment, (Label) null, opCode, args);
	}

	@Override
	public Void visitBlock(BlockContext ctx) {
		super.visitBlock(ctx);
		copyReg(ctx.expr(), ctx);
		return null;
	}

	@Override
	public Void visitDeclare(DeclareContext ctx) {
		return null;
	}

	@Override
	public Void visitParens(ParensContext ctx) {
		super.visitParens(ctx);
		copyReg(ctx.expr(), ctx);
		return null;
	}

	@Override
	public Void visitProgram(ProgramContext ctx) {
		emit("initialise temp ARP", OpCode.loadI, ZERO, tempArp);

		return visitChildren(ctx);
	}

	@Override
	public Void visitExpr(ExprContext ctx) {
		for (int i = 0; i < ctx.singleExpr().size() - 1; i++) {
			visit(ctx.singleExpr(i));
			freeReg(ctx.singleExpr(i));
		}
		ParserRuleContext last = ctx.singleExpr(ctx.singleExpr().size() - 1);
		visit(last);
		copyReg(last, ctx);
		return null;
	}

	@Override
	public Void visitAssign(AssignContext ctx) {
		visit(ctx.singleExpr());
		copyReg(ctx.singleExpr(), ctx);

		if (SimpleType.CHAR.equals(an.types.get(ctx)))
			emit("to " + ctx.variable().getText(), OpCode.cstoreAI, getReg(ctx), getVar(ctx.variable()),
					new Num(an.variables.get(ctx.variable()).getOffset()));
		else
			emit("to " + ctx.variable().getText(), OpCode.storeAI, getReg(ctx), getVar(ctx.variable()),
					new Num(an.variables.get(ctx.variable()).getOffset()));

		return null;
	}

	@Override
	public Void visitCall(CallContext ctx) {
		String base = "call " + ctx.variable().getText() + " - ";
		Variable<Type> function = an.variables.get(ctx.variable());
		TupleType parameters = (TupleType) ((FunctionType) function.getType()).getParameter();
		visit(ctx.variable());
		freeReg(ctx.variable());
		emit(base + "target address", OpCode.loadAI, getVar(ctx.variable()),
				new Num(an.variables.get(ctx.variable()).getOffset()), makeReg(ctx));
		emit(base + "target address", OpCode.push, getReg(ctx));
		freeReg(ctx);

		for (int i = 0; i < ctx.expr().size(); i++) {
			visit(ctx.expr(i));

			if (SimpleType.CHAR.equals(parameters.get(i)))
				emit(base + "param " + i, OpCode.cpush, getReg(ctx.expr(i)));
			else
				emit(base + "param " + i, OpCode.push, getReg(ctx.expr(i)));
			freeReg(ctx.expr(i));
		}

		emit(base + "set new ARP", OpCode.addI, arp,
				new Num(((FunctionType) an.function.get(ctx).getType()).getLocalDataSize() + ARBASESIZE), tempArp);

		makeReg(ctx);
		for (int i = ctx.expr().size() - 1; i >= 0; i--) {
			if (SimpleType.CHAR.equals(parameters.get(i))) {
				emit(base + "param " + i, OpCode.cpop, getReg(ctx));
				emit(base + "param " + i, OpCode.cstoreAI, getReg(ctx), tempArp, new Num(parameters.getOffset(i)));
			} else {
				emit(base + "param " + i, OpCode.pop, getReg(ctx));
				emit(base + "param " + i, OpCode.storeAI, getReg(ctx), tempArp, new Num(parameters.getOffset(i)));
			}
		}

		emit(base + "link caller ARP", OpCode.storeAI, arp, tempArp, OFFSET_ARP);

		emit(base + "travelling ALs", OpCode.i2i, arp, getReg(ctx));
		for (int i = an.variables.get(ctx.variable()).getDepth(); i <= an.function.get(ctx).getDepth(); i++)
			emit(base + StringUtils.repeat(' ', i) + "\\", OpCode.loadAI, getReg(ctx), OFFSET_AL, getReg(ctx));
		emit(base + "link caller AL", OpCode.storeAI, getReg(ctx), tempArp, OFFSET_AL);

		Label label = new Label("ret" + ctx.hashCode());
		emit(base + "load return address", OpCode.loadI, new Num(label), getReg(ctx));
		emit(base + "set return address", OpCode.storeAI, getReg(ctx), tempArp, OFFSET_RETURN_ADDR);
		emit(base + "move ARP", OpCode.i2i, tempArp, arp);
		emit(base + "target address", OpCode.pop, getReg(ctx));
		freeReg(ctx);

		Stack<Reg> inUse = new Stack<>();
		for (Reg reg : regInUse) {
			inUse.push(reg);
			emit(base + "register save " + reg.getName(), OpCode.push, reg);
		}

		emit(base + "execute", OpCode.jump, getReg(ctx));

		emit(base + "load result", label, OpCode.loadAI, arp, OFFSET_RETURN_VAL, makeReg(ctx));
		while (!inUse.isEmpty()) {
			Reg reg = inUse.pop();
			emit(base + "restoring register save " + reg.getName(), OpCode.pop, reg);
		}
		emit(base + "reset ARP", OpCode.loadAI, arp, OFFSET_ARP, arp);

		return null;
	}

	@Override
	public Void visitFuncDeclare(FuncDeclareContext ctx) {
		String base = "define " + ctx.IDENTIFIER(0).getText() + " - ";
		Label skip = new Label("s" + ctx.hashCode());
		emit(base + "jump over body", OpCode.jumpI, skip);
		int targetAddress = emit(base + "entry point", OpCode.nop).getLine();
		visit(ctx.body);
		emit(base + "move result", OpCode.storeAI, getReg(ctx.body), arp, OFFSET_RETURN_VAL);
		freeReg(ctx.body);
		emit(base + "load return address", OpCode.loadAI, arp, OFFSET_RETURN_ADDR, makeReg(ctx));
		emit(base + "go to return address", OpCode.jump, getReg(ctx));
		emit(base + "load target address", skip, OpCode.loadI, new Num(targetAddress), getReg(ctx));
		emit(base + "set target address", OpCode.storeAI, getReg(ctx), getVar(ctx),
				new Num(an.variables.get(ctx).getOffset()));
		freeReg(ctx);
		an.registers.removeFrom(ctx);
		return null;
	}

	@Override
	public Void visitBool(BoolContext ctx) {
		emit(ctx.getText(), OpCode.loadI, new Num(ctx.TRUE() != null ? 1 : 0), makeReg(ctx));
		return null;
	}

	@Override
	public Void visitNumber(NumberContext ctx) {
		emit(ctx.getText(), OpCode.loadI, new Num(Integer.parseInt(ctx.LITERAL10().getText())), makeReg(ctx));
		return null;
	}

	@Override
	public Void visitInfix1(Infix1Context ctx) {
		visitChildren(ctx);
		emit(ctx.getChild(1).getText(), ops.get(ctx.op.getType()), getReg(ctx.lhs), getReg(ctx.rhs),
				copyReg(ctx.lhs, ctx));
		freeReg(ctx.rhs);
		return null;
	}

	@Override
	public Void visitInfix2(Infix2Context ctx) {
		visitChildren(ctx);
		emit(ctx.getChild(1).getText(), ops.get(ctx.op.getType()), getReg(ctx.lhs), getReg(ctx.rhs),
				copyReg(ctx.lhs, ctx));
		freeReg(ctx.rhs);
		return null;
	}

	@Override
	public Void visitInfix3(Infix3Context ctx) {
		visitChildren(ctx);
		emit(ctx.getChild(1).getText(), ops.get(ctx.op.getType()), getReg(ctx.lhs), getReg(ctx.rhs),
				copyReg(ctx.lhs, ctx));
		freeReg(ctx.rhs);
		return null;
	}

	@Override
	public Void visitInfix4(Infix4Context ctx) {
		visitChildren(ctx);
		emit(ctx.getChild(1).getText(), ops.get(ctx.AND().getSymbol().getType()), getReg(ctx.lhs), getReg(ctx.rhs),
				copyReg(ctx.lhs, ctx));
		freeReg(ctx.rhs);
		return null;
	}

	@Override
	public Void visitInfix5(Infix5Context ctx) {
		visitChildren(ctx);
		emit(ctx.getChild(1).getText(), ops.get(ctx.OR().getSymbol().getType()), getReg(ctx.lhs), getReg(ctx.rhs),
				copyReg(ctx.lhs, ctx));
		freeReg(ctx.rhs);
		return null;
	}

	@Override
	public Void visitPrefix1(Prefix1Context ctx) {
		visitChildren(ctx);
		copyReg(ctx.singleExpr(), ctx);

		switch (ctx.op.getType()) {
		case BoppiLexer.MINUS:
			emit("unary -", OpCode.rsubI, getReg(ctx), ZERO, getReg(ctx));
		}
		return null;
	}

	@Override
	public Void visitVar(VarContext ctx) {
		if (SimpleType.CHAR.equals(an.types.get(ctx)))
			emit("var " + ctx.variable().getText(), OpCode.cloadAI, getVar(ctx.variable()),
					new Num(an.variables.get(ctx.variable()).getOffset()), makeReg(ctx));
		else
			emit("var " + ctx.variable().getText(), OpCode.loadAI, getVar(ctx.variable()),
					new Num(an.variables.get(ctx.variable()).getOffset()), makeReg(ctx));

		return null;
	}

	@Override
	public Void visitChar(CharContext ctx) {
		emit(ctx.getText(), OpCode.loadI, new Num(ctx.CHAR().getText().codePointAt(1)), makeReg(ctx));
		emit(ctx.getText(), OpCode.i2c, getReg(ctx), getReg(ctx));
		return null;
	}

	@Override
	public Void visitRead(ReadContext ctx) {
		if (ctx.variable().size() > 1) {
			for (VariableContext expr : ctx.variable()) {
				visit(expr);
				if (SimpleType.BOOL.equals(an.types.get(expr)) || SimpleType.INT.equals(an.types.get(expr))) {
					emit("", OpCode.in, new Str(""), makeReg(ctx));
					emit("var " + expr.getText(), OpCode.storeAI, getReg(ctx), getVar(expr),
							new Num(an.variables.get(expr).getOffset()));
				} else if (SimpleType.CHAR.equals(an.types.get(expr))) {
					// Get input until at least 1 character
					Label getTarget = new Label("lcin_l" + ctx.hashCode());
					Label continueTarget = new Label("lcin_e" + ctx.hashCode());

					Reg countReg = makeReg(ctx);
					makeReg(ctx);

					emit("", getTarget, OpCode.cin, new Str(""));
					emit("str length", OpCode.pop, countReg);
					emit("repeat if 0 length", OpCode.cbr, countReg, continueTarget, getTarget);

					// Get character
					emit("pop whole string", continueTarget, OpCode.cpop, getReg(ctx));
					emit("var " + expr.getText(), OpCode.cstoreAI, getReg(ctx), getVar(expr),
							new Num(an.variables.get(expr).getOffset()));

					// Pop all remaining characters
					Label loopTarget = new Label("lcpop_l" + ctx.hashCode());
					Label iterTarget = new Label("lcpop_c" + ctx.hashCode());
					Label stopTarget = new Label("lcpop_e" + ctx.hashCode());
					emit("pop remaining", loopTarget, OpCode.subI, countReg, new Num(1), countReg);
					emit("pop remaining", OpCode.cbr, countReg, iterTarget, stopTarget);
					emit("pop remaining", iterTarget, OpCode.cpop, getReg(ctx));
					emit("pop remaining", OpCode.jumpI, loopTarget);
					emit("pop remaining", stopTarget, OpCode.nop);
				} else {
					emit("????????", OpCode.out, new Str("Unknown type: "), copyReg(expr, ctx));
				}
				freeReg(ctx);
			}
		} else {
			VariableContext expr = ctx.variable(0);
			visit(expr);
			if (SimpleType.BOOL.equals(an.types.get(expr)) || SimpleType.INT.equals(an.types.get(expr))) {
				emit("", OpCode.in, new Str(""), makeReg(ctx));
				emit("var " + expr.getText(), OpCode.storeAI, getReg(ctx), getVar(expr),
						new Num(an.variables.get(expr).getOffset()));
			} else if (SimpleType.CHAR.equals(an.types.get(expr))) {
				// Get input until at least 1 character
				Label getTarget = new Label("lcin_l" + ctx.hashCode());
				Label continueTarget = new Label("lcin_e" + ctx.hashCode());

				Reg countReg = makeReg(ctx);
				makeReg(ctx);

				emit("", getTarget, OpCode.cin, new Str(""));
				emit("str length", OpCode.pop, countReg);
				emit("repeat if 0 length", OpCode.cbr, countReg, continueTarget, getTarget);

				// Get character
				emit("pop whole string", continueTarget, OpCode.cpop, getReg(ctx));
				emit("var " + expr.getText(), OpCode.cstoreAI, getReg(ctx), getVar(expr),
						new Num(an.variables.get(expr).getOffset()));

				// Pop all remaining characters
				Label loopTarget = new Label("lcpop_l" + ctx.hashCode());
				Label iterTarget = new Label("lcpop_c" + ctx.hashCode());
				Label stopTarget = new Label("lcpop_e" + ctx.hashCode());
				emit("pop remaining", loopTarget, OpCode.subI, countReg, new Num(1), countReg);
				emit("pop remaining", OpCode.cbr, countReg, iterTarget, stopTarget);
				emit("pop remaining", iterTarget, OpCode.cpop, getReg(ctx));
				emit("pop remaining", OpCode.jumpI, loopTarget);
				emit("pop remaining", stopTarget, OpCode.nop);
				// Reload value
				emit("var " + expr.getText(), OpCode.cloadAI, getVar(expr), new Num(an.variables.get(expr).getOffset()),
						getReg(ctx));
			}
		}
		return null;
	}

	@Override
	public Void visitIf(IfContext ctx) {
		Label toTrue = new Label("if_t" + ctx.hashCode());
		Label toFalse = new Label("if_f" + ctx.hashCode());
		Label toEnd = new Label("if_e" + ctx.hashCode());

		visit(ctx.cond);
		if (ctx.onFalse == null) {
			emit("", OpCode.cbr, getReg(ctx.cond), toTrue, toEnd);
			freeReg(ctx.cond);

			emit("", toTrue, OpCode.nop);
			visit(ctx.onTrue);
			freeReg(ctx.onTrue);
		} else {
			emit("", OpCode.cbr, getReg(ctx.cond), toTrue, toFalse);
			freeReg(ctx.cond);

			emit("", toTrue, OpCode.nop);
			visit(ctx.onTrue);
			freeReg(ctx.onTrue);
			if (an.types.get(ctx) != SimpleType.VOID)
				emit("result", OpCode.i2i, getReg(ctx.onTrue), makeReg(ctx));
			emit("", OpCode.jumpI, toEnd);
			freeReg(ctx);

			emit("", toFalse, OpCode.nop);
			visit(ctx.onFalse);
			freeReg(ctx.onFalse);
			if (an.types.get(ctx) != SimpleType.VOID)
				emit("result", OpCode.i2i, getReg(ctx.onFalse), makeReg(ctx));

		}
		emit("end target", toEnd, OpCode.nop);

		return null;
	}

	@Override
	public Void visitWhile(WhileContext ctx) {
		Label toLoop = new Label("while_t" + ctx.hashCode());
		Label toCond = new Label("while_f" + ctx.hashCode());
		Label toEnd = new Label("while_e" + ctx.hashCode());

		emit("to condition", OpCode.jumpI, toCond);
		emit("loop target", toLoop, OpCode.nop);
		visit(ctx.onTrue);
		freeReg(ctx.onTrue);
		emit("condition target", toCond, OpCode.nop);
		visit(ctx.cond);
		freeReg(ctx.cond);
		emit("", OpCode.cbr, getReg(ctx.cond), toLoop, toEnd);
		emit("end target", toEnd, OpCode.nop);

		return null;
	}

	@Override
	public Void visitWrite(WriteContext ctx) {
		if (ctx.expr().size() > 1) {
			for (ExprContext expr : ctx.expr()) {
				visit(expr);
				if (SimpleType.BOOL.equals(an.types.get(expr)) || SimpleType.INT.equals(an.types.get(expr))) {
					emit("", OpCode.out, new Str(""), getReg(expr));
				} else if (SimpleType.CHAR.equals(an.types.get(expr))) {
					makeReg(ctx);
					emit("push character", OpCode.cpush, getReg(expr));
					emit("load 1", OpCode.loadI, new Num(1), getReg(ctx));
					emit("push 1", OpCode.push, getReg(ctx));
					emit("print character", OpCode.cout, new Str(""));
					freeReg(ctx);
				} else {
					emit("??????????????", OpCode.out, new Str("Unknown type: "), getReg(expr));
				}
				freeReg(expr);
			}
			an.registers.removeFrom(ctx);
		} else {
			ExprContext expr = ctx.expr(0);
			visit(expr);
			if (SimpleType.BOOL.equals(an.types.get(expr)) || SimpleType.INT.equals(an.types.get(expr))) {
				emit("", OpCode.out, new Str(""), getReg(expr));
			} else if (SimpleType.CHAR.equals(an.types.get(expr))) {
				makeReg(ctx);
				emit("push character", OpCode.cpush, getReg(expr));
				emit("load 1", OpCode.loadI, new Num(1), getReg(ctx));
				emit("push 1", OpCode.push, getReg(ctx));
				emit("print character", OpCode.cout, new Str(""));
				freeReg(ctx);
			}

			copyReg(expr, ctx);
		}
		return null;
	}

	/**
	 * Returns an error message for a given parse tree node.
	 * 
	 * @param ctx
	 *            the parse tree node at which the error occurred
	 * @param message
	 *            the error message
	 * @param args
	 *            arguments for the message, see {@link String#format}
	 */
	private String getError(ParserRuleContext node, String message, Object... args) {
		int line = node.getStart().getLine();
		int column = node.getStart().getCharPositionInLine();
		return String.format("Line %d:%d - %s", line, column, String.format(message, args));
	}
}
