package pp.s1184725.boppi;

import java.util.*;
import java.util.function.Supplier;

/**
 * This class maps string identifiers to {@link Variable} instances. It takes
 * care of lexical scope, memory offsets and variable lifetime. This
 * implementation has a theoretical lookup time of O(1), scope opening time of
 * O(1) and scope closing time of O(k log k log n) with k the number of
 * identifiers going out of scope. It is agnostic to the type system used.
 *
 * @param <T>
 *            the typing class
 */
public class CachingSymbolTable<T extends Type> {
	protected Stack<Map<String, Variable<T>>> symbolMapStack;
	protected Stack<Integer> offsets;
	protected Stack<Integer> functionSizes;
	protected Map<String, Variable<T>> symbolCache;
	protected int offset;
	protected int functionDepth;
	protected int functionSize;

	/**
	 * Creates an empty symbol table with no open scope.
	 */
	public CachingSymbolTable() {
		symbolMapStack = new Stack<>();
		symbolCache = new HashMap<>();
		offsets = new Stack<>();
		functionSizes = new Stack<>();
		offset = 0;
		functionDepth = 0;
		functionSize = 0;
	}

	/**
	 * Opens a lexical scope.
	 */
	public void openScope() {
		symbolMapStack.push(new HashMap<>());
		offsets.push(offset);
	}

	/**
	 * Opens a lexical scope for a function declaration, resetting the local
	 * offset and increasing the lookup depth.
	 */
	public void openFunctionScope() {
		openScope();
		functionSizes.push(functionSize);
		functionDepth++;
		functionSize = 0;
		offset = 0;
	}

	/**
	 * Closes a lexical scope, removing all declarations within this scope.
	 * 
	 * @throws EmptyStackException
	 *             if no scope is open
	 */
	public void closeScope() throws EmptyStackException {
		Map<String, Variable<T>> deletions = symbolMapStack.pop();
		for (String identifier : deletions.keySet()) {
			Optional<Variable<T>> maybeType = symbolMapStack.stream().filter((map) -> map.containsKey(identifier))
					.map((map) -> map.get(identifier)).reduce((__, snd) -> snd);
			maybeType.ifPresent((var) -> symbolCache.replace(identifier, var));
			if (!maybeType.isPresent())
				symbolCache.remove(identifier);
		}
		offset = offsets.pop();
	}

	/**
	 * Closes a lexical scope for a function declaration, removing all
	 * declarations within this scope, restoring the local offset and decreasing
	 * the lookup depth.
	 * 
	 * @throws EmptyStackException
	 *             if no scope is open
	 */
	public void closeFunctionScope() throws EmptyStackException {
		closeScope();
		functionSize = functionSizes.pop();
		functionDepth--;
	}

	/**
	 * Opens a scope, executes the given code and closes the scope. This is a
	 * helper method to make sure calls to {@link #openScope()} and
	 * {@link #closeScope()} are balanced. {@code code} must take no arguments
	 * and must return a value.
	 * 
	 * @param <U>
	 *            the type of the code's return value
	 * @param code
	 *            the code to execute within the scope
	 * @return the return value of the code
	 */
	public <U> U withScope(Supplier<U> code) {
		openScope();
		U result = code.get();
		closeScope();
		return result;
	}

	/**
	 * Opens a function scope, executes the given code and closes it. This is a
	 * helper method to make sure calls to {@link #openFunctionScope()} and
	 * {@link #closeFunctionScope()} are balanced and returns the local data
	 * size (in bytes) of the function. {@code code} must take no arguments and
	 * must not return a value.
	 * 
	 * @param code
	 *            the code to execute within the scope
	 * @return the return value of the code
	 */
	public int withFunctionScope(Runnable code) {
		openFunctionScope();
		code.run();
		int result = functionSize;
		closeFunctionScope();
		return result;
	}

	/**
	 * Associates an identifier with a certain type in the current lexical scope
	 * and returns the newly made variable instance that belongs to this
	 * identifier. Throws an exception if the identifier is declared in this
	 * scope already.
	 * 
	 * @param id
	 *            the name of the identifier
	 * @param type
	 *            the type of the identifier
	 * @return the type of the identifier
	 * @throws Exception
	 *             if the identifier is declared in the current scope already
	 * @throws EmptyStackException
	 *             if no scope is open
	 */
	public Variable<T> put(String id, T type) throws Exception, EmptyStackException {
		if (symbolMapStack.peek().containsKey(id))
			throw new Exception(String.format("Identifier '%s' already declared in this scope.", id));

		Variable<T> var = new Variable<>(type, functionDepth, offset);
		offset += type.getSize();
		functionSize = Math.max(functionSize, offset);

		symbolMapStack.peek().put(id, var);
		symbolCache.put(id, var);
		return var;
	}

	/**
	 * Returns whether the given identifier is declared.
	 * 
	 * @param id
	 *            the name of the identifier
	 * @return true if the identifier has a declared type
	 * @throws EmptyStackException
	 *             if no scope is open
	 */
	public boolean has(String id) throws EmptyStackException {
		if (symbolMapStack.isEmpty())
			throw new EmptyStackException();

		return symbolCache.containsKey(id);
	}

	/**
	 * Retrieves the type of an identifier, if any.
	 * 
	 * @param id
	 *            the name of the identifier
	 * @return the type of the identifier
	 * @throws Exception
	 *             if the identifier is not declared
	 * @throws EmptyStackException
	 *             if no scope is open
	 */
	public Variable<T> get(String id) throws Exception, EmptyStackException {
		if (!this.has(id))
			throw new Exception(String.format("Identifier '%s' is undeclared.", id));

		return symbolCache.get(id);
	}

	/**
	 * Returns the size, in bytes, required to store all the variables declared
	 * in the current function scope. Note that this is the size calculated up
	 * and including the last variable declaration, so make sure to call it
	 * right before leaving a function scope.
	 * 
	 * @return the number of bytes required to store all variables local to a
	 *         function
	 */
	public int getLocalDataSize() {
		return functionSize;
	}

	/**
	 * Returns the current lexical scope/function depth.
	 * 
	 * @return the function depth
	 */
	public int getFunctionDepth() {
		return functionDepth;
	}
}
