package pp.s1184725.boppi;

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

import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.tree.ParseTree;

import pp.s1184725.boppi.antlr.*;
import pp.s1184725.boppi.antlr.BoppiParser.*;
import pp.s1184725.boppi.exception.*;
import pp.s1184725.boppi.type.*;

/**
 * This class performs type checking and variable assignment on a bare parse
 * tree.
 * 
 * @author Frank Wibbelink
 */
public class BoppiChecker extends BoppiBaseVisitor<Type> {
	private Annotations an;
	private Logger log;
	private boolean inLhs, inType;

	/**
	 * Checks and annotates a program. Problems are reported to the given
	 * logger.
	 * 
	 * @param ast
	 *            the program's AST
	 * @param logger
	 *            the logger object to write warnings and errors to
	 * @return the generated annotations
	 */
	public static Annotations checkProgram(ParseTree ast, Logger logger) {
		BoppiChecker checker = new BoppiChecker(logger, new Annotations());
		checker.visit(ast);
		return checker.an;
	}

	protected BoppiChecker(Logger logger, Annotations annotations) {
		an = annotations;
		log = logger;
		inLhs = false;
		inType = false;
	}

	/**
	 * Checks whether two types are compatible and log an error if they do not.
	 * 
	 * @param type1
	 *            the left hand or actual type
	 * @param type2
	 *            the right hand or desired type
	 * @param node
	 *            the context used for logging
	 */
	private void checkConstraint(Type type1, Type type2, ParserRuleContext node) {
		if (!type2.equals(type1))
			log.severe(getError(node, Messages.getString("BoppiChecker.0"), type1.toString(), type2.toString())); //$NON-NLS-1$
	}

	/**
	 * 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(Messages.getString("BoppiChecker.1"), line, column, String.format(message, args)); //$NON-NLS-1$
	}

	/**
	 * Returns whether the current context resides within a given rule context
	 * within a set root (may be {@code null}). Effectively traverses between
	 * {@code current} and {@code root} in the parent hierarchy looking for a
	 * {@code rule} class.
	 * 
	 * @param <T>
	 *            the rule class
	 * @param ctx
	 *            the node to start traversing
	 * @param root
	 *            the node to stop traversing
	 * @param rule
	 *            the ParserRuleContext class to look for
	 * @return whether a matching rule is found between the current and root
	 *         node
	 */
	private <T extends ParserRuleContext> boolean isWithinRule(ParserRuleContext ctx, ParserRuleContext root,
			Class<T> rule) {
		for (ParserRuleContext node = ctx; node != root && node != null; node = node.getParent())
			if (node.getClass().isAssignableFrom(rule))
				return true;

		return false;
	}

	/**
	 * Visits the parse tree, then records the returned type in the type
	 * annotations and returns it too.
	 * 
	 * @param tree
	 *            the tree to visit
	 * @return the evaluated type of the tree
	 */
	@Override
	public Type visit(ParseTree tree) {
		Type type = super.visit(tree);
		an.types.put(tree, type);
		return type;
	}

	@Override
	public Type visitAssign(AssignContext ctx) {
		Type expr = visit(ctx.assignStat() != null ? ctx.assignStat() : ctx.expr());
		inLhs = true;
		Type var = visit(ctx.variable());
		inLhs = false;
		checkConstraint(expr, var, ctx);
		return expr;
	}

	@Override
	public Type visitBlock(BlockContext ctx) {
		return an.symbols.withScope(() -> visit(ctx.stats()));
	}

	@Override
	public Type visitCall(CallContext ctx) {
		Type t = visit(ctx.variable());

		if (t instanceof FunctionType) {
			FunctionType type = (FunctionType) t;
			TupleType parameters = (TupleType) type.getParameter();
			if (parameters.size() != ctx.expr().size())
				log.severe(getError(ctx, Messages.getString("BoppiChecker.2"), parameters.size(), ctx.expr().size())); //$NON-NLS-1$

			for (int i = 0; i < ctx.expr().size() && i < parameters.size(); i++)
				checkConstraint(visit(ctx.expr(i)), parameters.get(i), ctx.expr(i));

			return type.getReturn();
		} else {
			log.severe(getError(ctx, Messages.getString("BoppiChecker.3"), ctx.variable().getText())); //$NON-NLS-1$
			return SimpleType.VOID;
		}
	}

	@Override
	public Type visitDeclare(DeclareContext ctx) {
		try {
			Variable<Type> var = an.symbols.put(ctx.IDENTIFIER().getText(), visit(ctx.type()));
			var.setConstant(ctx.CONSTANT() != null);
			an.variables.put(ctx, var);
			an.variableRoot.put(var, ctx.getParent().getParent());

			if (var.getType() instanceof TupleType)
				log.severe(getError(ctx, Messages.getString("BoppiChecker.4"), //$NON-NLS-1$
						BoppiLexer.VOCABULARY.getLiteralName(BoppiLexer.BOOLTYPE),
						BoppiLexer.VOCABULARY.getLiteralName(BoppiLexer.CHARTYPE),
						BoppiLexer.VOCABULARY.getLiteralName(BoppiLexer.INTTYPE)));
		} catch (SymbolTableException e) {
			log.severe(getError(ctx, e.getMessage()));
		}
		return SimpleType.VOID;
	}

	@Override
	public Type visitDeclareFunction(DeclareFunctionContext ctx) {
		try {
			TupleType parameterTypes;
			if (ctx.parameters() != null)
				parameterTypes = new TupleType(
						ctx.parameters().type().stream().map(this::visit).collect(Collectors.toList()));
			else
				parameterTypes = TupleType.UNIT;

			Type returnType = ctx.result != null ? visit(ctx.result) : SimpleType.VOID;
			FunctionType type = new FunctionType(returnType, parameterTypes);

			Variable<Type> func = an.symbols.put(ctx.name.getText(), type);
			func.setConstant(true);
			func.assign();
			an.variables.put(ctx, func);

			an.function.put(ctx, an.symbols.withFunctionScope(() -> {
				if (ctx.parameters() != null)
					visit(ctx.parameters());

				Type resultType = an.symbols.withScope(() -> visit(ctx.body));

				if (returnType != SimpleType.VOID)
					checkConstraint(resultType, returnType, ctx);
			}));
		} catch (SymbolTableException e) {
			log.severe(getError(ctx, e.getMessage()));
		}

		return SimpleType.VOID;
	}

	@Override
	public Type visitDefineArray(DefineArrayContext ctx) {
		checkConstraint(visit(ctx.size), SimpleType.INT, ctx);

		return new ArrayType(visit(ctx.type()));
	}

	@Override
	public Type visitGetVariable(GetVariableContext ctx) {
		return visit(ctx.variable());
	}

	@Override
	public Type visitIf(IfContext ctx) {
		return an.symbols.withScope(() -> {
			checkConstraint(visit(ctx.cond), SimpleType.BOOL, ctx.cond);

			if (ctx.onFalse != null) {
				Type trueType = an.symbols.withScope(() -> visit(ctx.onTrue));
				Type falseType = an.symbols.withScope(() -> visit(ctx.onFalse));

				return (trueType.equals(falseType)) ? trueType : SimpleType.VOID;
			} else {
				an.symbols.withScope(() -> visit(ctx.onTrue));
				return SimpleType.VOID;
			}
		});
	}

	@Override
	public Type visitInfix1(Infix1Context ctx) {
		checkConstraint(visit(ctx.lhs), SimpleType.INT, ctx.lhs);
		checkConstraint(visit(ctx.rhs), SimpleType.INT, ctx.rhs);

		return SimpleType.INT;
	}

	@Override
	public Type visitInfix2(Infix2Context ctx) {
		checkConstraint(visit(ctx.lhs), SimpleType.INT, ctx.lhs);
		checkConstraint(visit(ctx.rhs), SimpleType.INT, ctx.rhs);

		return SimpleType.INT;
	}

	@Override
	public Type visitInfix3(Infix3Context ctx) {
		Type lht = visit(ctx.lhs);
		Type rht = visit(ctx.rhs);
		checkConstraint(lht, rht, ctx);

		switch (ctx.op.getType()) {
		case BoppiLexer.LT:
		case BoppiLexer.LEQ:
		case BoppiLexer.GTE:
		case BoppiLexer.GT:
			checkConstraint(lht, SimpleType.INT, ctx);
		}

		return SimpleType.BOOL;
	}

	@Override
	public Type visitInfix4(Infix4Context ctx) {
		checkConstraint(visit(ctx.lhs), SimpleType.BOOL, ctx.lhs);
		checkConstraint(visit(ctx.rhs), SimpleType.BOOL, ctx.rhs);

		return SimpleType.BOOL;
	}

	@Override
	public Type visitInfix5(Infix5Context ctx) {
		checkConstraint(visit(ctx.lhs), SimpleType.BOOL, ctx.lhs);
		checkConstraint(visit(ctx.rhs), SimpleType.BOOL, ctx.rhs);

		return SimpleType.BOOL;
	}

	@Override
	public Type visitLiteralArray(LiteralArrayContext ctx) {
		Type elementType = visit(ctx.expr(0));

		for (int i = 1; i < ctx.expr().size(); i++)
			checkConstraint(visit(ctx.expr(i)), elementType, ctx);

		if (SimpleType.VOID.equals(elementType))
			log.severe(getError(ctx, Messages.getString("BoppiChecker.16"))); //$NON-NLS-1$

		return new ArrayType(elementType);
	}

	@Override
	public Type visitLiteralBoolean(LiteralBooleanContext ctx) {
		return SimpleType.BOOL;
	}

	@Override
	public Type visitLiteralCharacter(LiteralCharacterContext ctx) {
		return SimpleType.CHAR;
	}

	@Override
	public Type visitLiteralInteger(LiteralIntegerContext ctx) {
		return SimpleType.INT;
	}

	@Override
	public Type visitParens(ParensContext ctx) {
		return visit(ctx.stats());
	}

	@Override
	public Type visitPrefix1(Prefix1Context ctx) {
		Type type = visit(ctx.expr());

		switch (ctx.op.getType()) {
		case BoppiLexer.NOT:
			checkConstraint(type, SimpleType.BOOL, ctx.expr());
			break;

		case BoppiLexer.PLUS:
		case BoppiLexer.MINUS:
			checkConstraint(type, SimpleType.INT, ctx.expr());
			break;
		}

		return type;
	}

	@Override
	public Type visitParameters(ParametersContext ctx) {
		for (int i = 0; i < ctx.type().size(); i++)
			try {
				an.symbols.put(ctx.IDENTIFIER(i).getText(), an.types.get(ctx.type(i))).assign();
			} catch (SymbolTableException e) {
				log.severe(getError(ctx, e.getMessage()));
			}
		return null;
	}

	@Override
	public Type visitProgram(ProgramContext ctx) {
		an.function.put(ctx, an.symbols.withFunctionScope(() -> super.visitProgram(ctx)));
		return null;
	}

	@Override
	public Type visitRead(ReadContext ctx) {
		Type result;

		inLhs = true;
		if (ctx.variable().size() == 1) {
			result = visit(ctx.variable(0));
		} else {
			ctx.variable().forEach(this::visit);
			result = SimpleType.VOID;
		}
		inLhs = false;

		return result;
	}

	@Override
	public Type visitStat(StatContext ctx) {
		return visit(ctx.getChild(0));
	}

	@Override
	public Type visitStats(StatsContext ctx) {
		Type resultType = ctx.stat().stream().map(this::visit).reduce((__, snd) -> snd).get();

		ParseTree lastExpr = ctx.stat(ctx.stat().size() - 1).getChild(0);
		if (lastExpr instanceof DeclareContext || lastExpr instanceof DeclareFunctionContext)
			log.severe(getError(ctx, Messages.getString("BoppiChecker.5"))); //$NON-NLS-1$

		return resultType;
	}

	@Override
	public Type visitTypeArray(TypeArrayContext ctx) {
		return new ArrayType(visit(ctx.type()));
	}

	@Override
	public Type visitTypeFunction(TypeFunctionContext ctx) {
		return new FunctionType(visit(ctx.type(1)), visit(ctx.type(0)));
	}

	@Override
	public Type visitTypeSimple(TypeSimpleContext ctx) {
		return SimpleType.fromToken(ctx.staticType.getType());
	}

	@Override
	public Type visitTypeTuple(TypeTupleContext ctx) {
		return new TupleType(ctx.type().stream().map(this::visit).collect(Collectors.toList()));
	}

	@Override
	public Type visitTypeVariable(TypeVariableContext ctx) {
		inType = true;
		Type type = visit(ctx.variable());
		inType = false;
		return type;
	}

	@Override
	public Type visitVariableArray(VariableArrayContext ctx) {
		boolean wasOnLhs = inLhs;
		inLhs = false;
		checkConstraint(visit(ctx.expr()), SimpleType.INT, ctx);
		inLhs = wasOnLhs;

		Type t = visit(ctx.variable());

		if (t instanceof ArrayType) {
			ArrayType arrayType = (ArrayType) t;
			return arrayType.getType();
		} else {
			log.severe(getError(ctx, Messages.getString("BoppiChecker.7"), t)); //$NON-NLS-1$
			return SimpleType.VOID;
		}
	}

	@Override
	public Type visitVariableProperty(VariablePropertyContext ctx) {
		Type varType = visit(ctx.variable());
		String prop = ctx.IDENTIFIER().getText();

		if (varType instanceof ArrayType) {
			if (prop.equals("length")) { //$NON-NLS-1$
				return SimpleType.INT;
			} else {
				log.severe(getError(ctx, Messages.getString("BoppiChecker.9"), prop)); //$NON-NLS-1$
				return SimpleType.VOID;
			}
		} else {
			log.severe(getError(ctx, Messages.getString("BoppiChecker.10"))); //$NON-NLS-1$
			return SimpleType.VOID;
		}
	}

	@Override
	public Type visitVariableSimple(VariableSimpleContext ctx) {
		try {
			Variable<Type> var = an.symbols.get(ctx.getText());

			if (!inType) {
				if (inLhs) {
					if (var.isConstant()) {
						if (var.isAssigned())
							log.severe(getError(ctx, Messages.getString("BoppiChecker.8"), ctx.getText())); //$NON-NLS-1$
						else if (var.isPossiblyAssigned())
							log.warning(getError(ctx, Messages.getString("BoppiChecker.13"), ctx.getText())); //$NON-NLS-1$
					}

					if (inLhs) {
						if (isWithinRule(ctx, an.variableRoot.get(var), IfContext.class)
								|| isWithinRule(ctx, an.variableRoot.get(var), WhileContext.class)
								|| isWithinRule(ctx, an.variableRoot.get(var), DeclareFunctionContext.class))
							var.possiblyAssign();
						else
							var.assign();
					}
				} else {
					if (!var.isAssigned()) {
						if (var.isPossiblyAssigned()
								|| isWithinRule(ctx, an.variableRoot.get(var), DeclareFunctionContext.class))
							log.warning(getError(ctx, Messages.getString("BoppiChecker.14"), ctx.getText())); //$NON-NLS-1$
						else
							log.severe(getError(ctx, Messages.getString("BoppiChecker.15"), ctx.getText())); //$NON-NLS-1$
					}
				}
			}

			try {
				an.function.put(ctx, an.symbols.getFunctionScope());
			} catch (NoProgramException e) {
				log.severe(getError(ctx, e.getMessage()));
			}

			an.variables.put(ctx, var);
			return var.getType();
		} catch (SymbolTableException e) {
			log.severe(getError(ctx, e.getMessage()));
			return SimpleType.VOID;
		}
	}

	@Override
	public Type visitWhile(WhileContext ctx) {
		return an.symbols.withScope(() -> {
			checkConstraint(visit(ctx.cond), SimpleType.BOOL, ctx.cond);
			an.symbols.withScope(() -> visit(ctx.onTrue));
			return SimpleType.VOID;
		});
	}

	@Override
	public Type visitWrite(WriteContext ctx) {
		if (ctx.expr().size() == 1) {
			Type type = visit(ctx.expr(0));
			if (SimpleType.VOID.equals(type))
				log.severe(getError(ctx, Messages.getString("BoppiChecker.11"), type)); //$NON-NLS-1$

			return type;
		} else {
			ctx.expr().stream().map(this::visit).forEach((type) -> {
				if (SimpleType.VOID.equals(type))
					log.severe(getError(ctx, Messages.getString("BoppiChecker.12"), type)); //$NON-NLS-1$
			});
			return SimpleType.VOID;
		}
	}
}
