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.ParseTree;

import pp.s1184725.boppi.antlr.BoppiBaseVisitor;
import pp.s1184725.boppi.antlr.BoppiLexer;
import pp.s1184725.boppi.antlr.BoppiParser.*;

/**
 * This class performs type checking and variable assignment on a bare parse
 * tree.
 *
 */
public class BoppiChecker extends BoppiBaseVisitor<Type> {
	private CachingSymbolTable<Type> symbols;
	private Annotations an;
	private Logger log;

	/**
	 * 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);
		checker.visit(ast);
		return checker.an;
	}

	protected BoppiChecker(Logger logger) {
		symbols = new CachingSymbolTable<>();
		an = new Annotations();
		log = logger;
	}

	@Override
	public Type visit(ParseTree tree) {
		Type type = super.visit(tree);
		an.types.put(tree, type);
		if (an.function.get(tree) == null)
			an.function.put(tree, an.currentFunction.peek());
		return type;
	}

	@Override
	public Type visitAssign(AssignContext ctx) {
		Type expr = visit(ctx.singleExpr());
		Type var = visit(ctx.variable());
		checkConstraint(expr, var, ctx);
		return expr;
	}

	@Override
	public Type visitBool(BoolContext ctx) {
		return SimpleType.BOOL;
	}

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

	@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, "Expected %d arguments but got %d.", parameters.size(), ctx.expr().size()));

			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, "'%s' is not a function.", ctx.variable().getText()));
			return SimpleType.VOID;
		}
	}

	@Override
	public Type visitChar(CharContext ctx) {
		return SimpleType.CHAR;
	}

	@Override
	public Type visitExpr(ExprContext ctx) {
		if (ctx.singleExpr(ctx.singleExpr().size() - 1) instanceof DeclareContext)
			log.severe(getError(ctx, "Compound expression ends with declaration."));

		return ctx.singleExpr().stream().map(this::visit).reduce((__, snd) -> snd).get();
	}

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

			if (var.getType() instanceof TupleType)
				log.severe(getError(ctx, "Variable must have a type %s, %s, %s or function.",
						BoppiLexer.VOCABULARY.getLiteralName(BoppiLexer.BOOLTYPE),
						BoppiLexer.VOCABULARY.getLiteralName(BoppiLexer.CHARTYPE),
						BoppiLexer.VOCABULARY.getLiteralName(BoppiLexer.INTTYPE)));

			if (!(var.getType() instanceof SimpleType))
				log.warning("Be careful only to pass pure functions outside their scope.");
		} catch (EmptyStackException e) {
			log.severe(getError(ctx, "Declaring variable outside a program."));
		} catch (Exception e) {
			log.severe(getError(ctx, e.getMessage()));
		}
		return SimpleType.VOID;
	}

	@Override
	public Type visitFuncDeclare(FuncDeclareContext ctx) {
		try {
			TupleType parameterTypes = new TupleType(
					ctx.type().stream().skip(1).map(this::visit).collect(Collectors.toList()));
			FunctionType type = new FunctionType(visit(ctx.result), parameterTypes);
			Variable<Type> func = symbols.put(ctx.name.getText(), type);
			an.variables.put(ctx, func);
			an.currentFunction.push(func);

			type.setLocalDataSize(symbols.withFunctionScope(() -> {
				for (int i = 1; i < ctx.type().size(); i++)
					try {
						symbols.put(ctx.IDENTIFIER(i).getText(), an.types.get(ctx.type(i)));
					} catch (Exception e) {
						log.severe(getError(ctx, e.getMessage()));
					}
				checkConstraint(symbols.withScope(() -> visit(ctx.body)), an.types.get(ctx.result), ctx);
			}));
		} catch (EmptyStackException e) {
			log.severe(getError(ctx, "Declaring function outside a program."));
		} catch (Exception e) {
			log.severe(getError(ctx, e.getMessage()));
		}
		an.currentFunction.pop();
		return SimpleType.VOID;
	}

	@Override
	public Type visitIf(IfContext ctx) {
		return symbols.withScope(() -> {
			visit(ctx.cond);
			if (ctx.onFalse != null) {
				Type trueType = symbols.withScope(() -> visit(ctx.onTrue));
				Type falseType = symbols.withScope(() -> visit(ctx.onFalse));

				return (trueType.equals(falseType)) ? trueType : SimpleType.VOID;
			} else {
				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 visitNumber(NumberContext ctx) {
		return SimpleType.INT;
	}

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

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

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

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

		return type;
	}

	@Override
	public Type visitProgram(ProgramContext ctx) {
		FunctionType main = new FunctionType(TupleType.UNIT, TupleType.UNIT);
		an.currentFunction.push(new Variable<Type>(main, 0, 0));
		main.setLocalDataSize(symbols.withFunctionScope(() -> super.visitProgram(ctx)));
		return null;
	}

	@Override
	public Type visitRead(ReadContext ctx) {
		if (ctx.variable().size() == 1) {
			return visit(ctx.variable(0));
		} else {
			ctx.variable().forEach(this::visit);
			return SimpleType.VOID;
		}
	}

	@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.children.stream().map(this::visit).collect(Collectors.toList()));
	}

	@Override
	public Type visitTypeVariable(TypeVariableContext ctx) {
		return visit(ctx.variable());
	}

	@Override
	public Type visitVar(VarContext ctx) {
		return visit(ctx.variable());
	}

	@Override
	public Type visitVariable(VariableContext ctx) {
		try {
			Variable<Type> var = symbols.get(ctx.getText());
			an.variables.put(ctx, var);
			return var.getType();
		} catch (EmptyStackException e) {
			log.severe(getError(ctx, "Using variable outside a program."));
		} catch (Exception e) {
			log.severe(getError(ctx, e.getMessage() != null ? e.getMessage() : "unknown error"));
		}

		return SimpleType.VOID;
	}

	@Override
	public Type visitWhile(WhileContext ctx) {
		return symbols.withScope(() -> {
			checkConstraint(visit(ctx.cond), SimpleType.BOOL, ctx.cond);
			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, "Cannot print argument of type %s.", type));

			return type;
		} else {
			ctx.expr().stream().map(this::visit).forEach((type) -> {
				if (SimpleType.VOID.equals(type))
					log.severe(getError(ctx, "Cannot print argument of type %s.", type));
			});
			return SimpleType.VOID;
		}
	}

	/**
	 * 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, "Could not match type %s with type %s.", type1.toString(), type2.toString()));
	}

	/**
	 * 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));
	}
}
