package pp.s1184725.boppi;

import java.util.*;
import java.util.function.*;
import java.util.logging.Logger;

import pp.iloc.model.Reg;

/**
 * Simple class to manage registers.
 */
public class RegisterPool {
	/**
	 * Recommended number of registers to use before throwing a warning.
	 */
	public static int RECOMMENDED_REGISTER_COUNT = 10;

	/**
	 * Register that must have a value of 0.
	 */
	public static final Reg ZERO = new Reg("r_nul");

	private Logger logger;
	private List<Reg> regFree, regInUse;

	/**
	 * Creates a new register pool
	 *
	 * @param logger
	 *            the logger to use
	 */
	public RegisterPool(Logger logger) {
		regFree = new ArrayList<>();
		regInUse = new ArrayList<>();
		this.logger = logger;
	}

	private void blockReg(Reg reg) {
		if (reg == null)
			logger.severe("INTERNAL: cannot reserve null register.");
		else if (!regFree.contains(reg)) {
			logger.severe(String.format("INTERNAL: cannot reserve register %s because it is in use.", reg.getName()));
		} else {
			regFree.remove(reg);
			regInUse.add(reg);
		}
	}

	/**
	 * Blocks two registers while running some code.
	 *
	 * @param <T>
	 *            the return type of the code
	 * @param r1
	 *            the first register to reserve
	 * @param r2
	 *            the second register to reserve
	 * @param code
	 *            the code to run
	 * @return the return value of the code
	 */
	public <T> T blockReg(Reg r1, Reg r2, Supplier<T> code) {
		return blockReg(r1, () -> blockReg(r2, code));
	}

	/**
	 * Blocks a register while running some code.
	 *
	 * @param r1
	 *            the register to reserve
	 * @param code
	 *            the code to run
	 */
	public void blockReg(Reg r1, Runnable code) {
		blockReg(r1);
		code.run();
		freeReg(r1);
	}

	/**
	 * Blocks a register while running some code.
	 *
	 * @param <T>
	 *            the return type of the code
	 * @param r1
	 *            the register to reserve
	 * @param code
	 *            the code to run
	 * @return the return value of the code
	 */
	public <T> T blockReg(Reg r1, Supplier<T> code) {
		blockReg(r1);
		T result = code.get();
		freeReg(r1);
		return result;
	}

	private void freeReg(Reg reg) {
		if (reg == null) {
			logger.severe("INTERNAL: cannot free null register.");
		} else if (!regInUse.contains(reg)) {
			logger.severe(String.format("INTERNAL: cannot free register %s because it is not in use.", reg.getName()));
		} else {
			regInUse.remove(reg);
			regFree.add(reg);
		}
	}

	/**
	 * Returns the list of registers currently in use.
	 *
	 * @return the list of registers in use
	 */
	public List<Reg> getInUse() {
		return regInUse;
	}

	private Reg makeReg() {
		return makeRegThatIsNot(null);
	}

	private Reg makeRegThatIsNot(Reg r1) {
		Reg reg = regFree.parallelStream().filter((r2) -> !r2.equals(r1)).findAny().orElseGet(() -> {
			if (regInUse.size() == RECOMMENDED_REGISTER_COUNT + 1)
				logger.warning(String.format("Using more than %d registers. Consider rebalancing your expressions.",
						RECOMMENDED_REGISTER_COUNT));

			return new Reg("r_" + (regInUse.size() + 1));
		});

		regFree.remove(reg);
		regInUse.add(reg);

		return reg;
	}

	/**
	 * Runs some code with two registers allocated.
	 *
	 * @param code
	 *            the code to run
	 * @return the second register used in the code
	 */
	public Reg withReg(BiConsumer<Reg, Reg> code) {
		return withReg((r2) -> withReg((r1) -> code.accept(r1, r2)));
	}

	/**
	 * Runs some code with a single register allocated.
	 *
	 * @param code
	 *            the code to run
	 * @return the register used in the code
	 */
	public Reg withReg(Consumer<Reg> code) {
		Reg r = makeReg();
		code.accept(r);
		freeReg(r);
		return r;
	}

}
