use core:lang;
use lang:bnf;
use lang:bs;
use lang:bs:macro;
use core:asm;
use core:io;

/**
 * A backtracking recursive descent parser.
 *
 * It behaves similarly to the recursive LL(1) parser in that it generates functions for each
 * production. The difference is that it tries matches until one in found, and backtracks on
 * failures. This means that it supports a larger set of languages, but side-effects from partial
 * parses are visible in calls to functions, and performance might suffer if backtracking is
 * required often. To resolve conflicts, the parser requires that productions have a priority set.
 *
 * The following structure is generated for each rule:
 *
 * T? Rx(State state, ...) {
 *     var saved = state.pos;
 *     if (a = match()) {
 *         if (b = match()) {
 *             // ...
 *             return match;
 *         }
 *     }
 *     state.pos = saved;
 *     if (a = match()) {
 *         // ...
 *         return match;
 *     }
 *     return null;
 * }
 */
class BacktrackingParser extends Parser {
	init(Scope env, SStr name) {
		init(env, name) {}
	}

	FnBody createBody(BSRawFn fn, Bool binary, Scope scope, Rule start, NameSet[] include) : override {
		Grammar grammar(start, include, true);

		{
			var recursive = grammar.leftRecursive;
			if (recursive.any) {
				StrBuf msg;
				msg << "Left recursive grammars are not supported by the backtracking parser. ";
				msg << "The grammar has at least the following left recursive components:\n";
				msg << join(recursive, "\n");
				throw GrammarError(name.pos, msg.toS);
			}
		}

		FnBody body(fn, scope);

		// Note: Parameter 0 is the input, parameter 1 is the start position.
		Var stateVar = {
			Actuals params;
			params.add(LocalVarAccess(SrcPos(), body.parameters[0]));
			params.add(LocalVarAccess(SrcPos(), body.parameters[1]));
			Var(body, stateType(binary), SStr("state"), params);
		};
		body.add(stateVar);

		// Create all required functions.
		FnCache cache(grammar, fn, scope, binary);
		cache.createAll();

		// Call the function for the start production.
		Actuals actuals;
		actuals.add(LocalVarAccess(SrcPos(), stateVar.var));
		for (Nat i = 2; i < body.parameters.count; i++)
			actuals.add(LocalVarAccess(SrcPos(), body.parameters[i]));

		FnCall call(fn.pos, scope, cache.get(grammar.start), actuals);

		Var resultVar(body, fn.result.type, SStr("result"), Actuals());
		body.add(resultVar);

		body.add(pattern(body) {
					result.value = ${call};
					result.end = state.pos;

					if (state.errors.any) {
						// Did we manage to parse everything?
						if (state.pos + 1 == state.pos) {
							// No error.
						} else {
							result.error = createError(state.errorMsg, state.longest);
						}
					}

					return result;
				});

		body;
	}

	/**
	 * State when parsing strings.
	 */
	class StrState extends RecursiveStrState {
		// List of errors. We make a string of them on demand.
		Str[] errors;

		// Create.
		init(Str input, Str:Iter start) {
			init(input, start) { longest = start; }
		}

		// Longest match.
		Str:Iter longest;

		// Create an error message.
		Str errorMsg() {
			StrBuf out;
			out << join(errors, "\n");
			out.toS;
		}

		// Set error.
		void error(Str error) : override {
			if (pos >= longest) {
				errors << error;
			}
		}

		// Update the position, set 'longest' if necessary.
		void update(Str:Iter p) {
			pos = p;
			if (p > longest) {
				longest = p;
				errors.clear();
			}
		}
	}

	/**
	 * State when parsing buffers.
	 */
	class BinaryState extends RecursiveBinaryState {
		Str error;

		init(Buffer input, Nat start) {
			init(input, start) { longest = start; }
		}

		// Longest match.
		Nat longest;

		// Set error.
		void error(Str error) : override {
			if (pos >= longest) {
				if (this.error.empty)
					this.error = error;
				else
					this.error = this.error + "\n" + error;
			}
		}

		// Update the position, set 'longest' if necessary.
		void update(Nat p) {
			pos = p;
			if (p > longest) {
				longest = p;
				error = "";
			}
		}
	}

	// Get the type of the state required, based on if we're in binary mode or not.
	Type stateType(Bool binary) : static {
		if (binary)
			named{BinaryState};
		else
			named{StrState};
	}

	/**
	 * Keep track of created functions.
	 */
	class FnCache on Compiler {
		init(Grammar grammar, Function parent, Scope scope, Bool binary) {
			init {
				grammar = grammar;
				parent = parent;
				scope = scope;
				binary = binary;
				thread = parent.declaredThread;
			}

			loadSpecial();
		}

		// Parent entity.
		Named parent;

		// Scope.
		Scope scope;

		// The grammar.
		Grammar grammar;

		// Running in binary mode?
		Bool binary;

		// Named thread to run on.
		NamedThread? thread;

		// Created functions. Not necessarily entirely done.
		private Rule->Function ruleFns;

		// Get a function. Does not populate it.
		Function get(Rule r) {
			if (fn = ruleFns.at(r))
				return fn;

			var params = r.params.clone;
			params.insert(0, ValParam(stateType(binary), "@state"));

			BSTreeFn fn(resultType(r.result), SStr(r.name), params, thread);
			ruleFns.put(r, fn);

			fn.parentLookup = parent;
			fn;
		}

		// Create all functions in the grammar.
		void createAll() {
			for (k, v in grammar.rules)
				create(k, v);
		}

		/**
		 * State common to the generator functions below.
		 */
		class ProdState on Compiler {
			init(SrcPos pos, ExprBlock block, Production production, Expr state, Expr failResult) {
				init {
					pos = pos;
					block = block;
					state = state;
					production = production;
					failResult = failResult;
				}
			}

			// For error messages.
			SrcPos pos;

			// The block we are currently generating code into.
			ExprBlock block;

			// Expression to access the state object passed to the parse function.
			Expr state;

			// Current production.
			Production production;

			// What to return on failure.
			Expr failResult;

			// Variables created so far.
			Str->LocalVar variables;

			// Start of captured segment. Populated when we encounter it.
			LocalVar? captureStart;

			// If in a repeat: the loop we are populating.
			Expr? repeatRoot;

			// If in a repeat: the block in the loop we are populating.
			ExprBlock? repeatBlock;

			// If in a repeat: the block used outside the loop.
			ExprBlock? repeatPrev;

			// Operations to "commit" whenever an repetation is complete.
			Expr[] repeatCommit;

			// Is the "me" variable created?
			Bool createdMeVar;

			// Check if "name" refers to a simple expression we are aware of (integers, booleans).
			Expr? simpleExpr(Str name) {
				if (v = name.long) {
					return NumLiteral(SrcPos(), v);
				} else if (name == "false") {
					return BoolLiteral(SrcPos(), false);
				} else if (name == "true") {
					return BoolLiteral(SrcPos(), true);
				} else {
					return null;
				}
			}

			// Get a variable. Throws appropriate message on error.
			Expr variable(Str name) {
				if (!variables.has(name)) {
					if (name == "me") {
						if (!createMeVar())
							throw GrammarError(pos, "The variable 'me' is not declared in this production, or evaluates to 'void'.");
					} else if (e = simpleExpr(name)) {
						return e;
					} else {
						StrBuf msg;
						msg << "The parameter '" << name << "' is used before it is declared. It needs to be ";
						msg << "initialized by one of the tokens to the left of the one where it is used.";
						throw GrammarError(pos, msg.toS);
					}
				}

				LocalVarAccess(pos, variables.get(name));
			}

			// Create the 'me' variable if possible.
			Bool createMeVar() {
				unless (result = production.result)
					return false;

				if (createdMeVar)
					return false;
				createdMeVar = true;

				// Initialize outside of any loops.
				ExprBlock initIn = block;
				if (repeatPrev)
					initIn = repeatPrev;

				if (params = production.resultParams) {
					// It is a function call!
					Actuals actuals;
					for (x in params) {
						if (variables.has(x)) {
							actuals.add(LocalVarAccess(pos, variables.get(x)));
						} else if (e = simpleExpr(x)) {
							actuals.add(e);
						} else {
							throw GrammarError(pos, "When initializing 'me': can not find a variable named ${x}.");
						}
					}

					// Note: This is technically done in the wrong context. Should be done in the
					// context of the production, not that of this function.
					Expr init = namedExpr(initIn, pos, result, actuals);
					if (init.result.type == Value()) {
						// It evaluates to 'void'. Make sure to still evaluate it.
						initIn.add(init);
						return false;
					}

					Var v(initIn, SStr("me"), init);
					initIn.add(v);
					variables.put("me", v.var);
				} else {
					// It is a variable! That means we can simply "rename" a local variable.
					if (result.count != 1)
						throw GrammarError(pos, "Unknown variable: ${result}");
					Str name = result[0].name;
					if (!variables.has(name))
						throw GrammarError(pos, "When initializing 'me': can not find a variable with the name ${name}");
					variables.put("me", variables.get(name));
				}
				true;
			}

			// Start a repeating section.
			void onRepStart() {
				if (production.repType == RepType:repNone)
					return;

				repeatPrev = block;

				Loop loop(SrcPos(), block);
				repeatRoot = loop;

				block = ExprBlock(SrcPos(), loop);
				repeatBlock = block;
				loop.doBody(block);

				Var posVar(block, SStr("@repPos"), namedExpr(block, SStr("pos"), state));
				block.add(posVar);
			}

			// End a repeating section.
			void onRepEnd() {
				unless (repeatRoot)
					return;
				unless (repeatPrev)
					return;
				unless (repeatBlock)
					return;

				// Check what to do if we failed. Regardless, we first have to restore the position.
				unless (oldPos = repeatBlock.variable(SimplePart("@repPos")))
					throw InternalError("Could not find the start-of-loop variable!");

				var posMember = namedExpr(repeatBlock, SStr("pos"), state);
				var assignOp = assignOperator(SStr("="), 0);
				repeatBlock.add(Operator(repeatBlock, posMember, assignOp, LocalVarAccess(SrcPos(), oldPos)));

				if (production.repType.skippable) {
					repeatBlock.add(Break(SrcPos(), repeatBlock));
				} else {
					// If we need one or more iteration, we need to ensure that the loop is executed
					// at least once!
					Var onceVar(repeatPrev, SStr("@once"), BoolLiteral(SrcPos(), false));
					repeatPrev.add(onceVar);

					LocalVarAccess once(SrcPos(), onceVar.var);

					// Set once to true when done.
					block.add(Operator(block, once, assignOperator(SStr("="), 0), BoolLiteral(SrcPos(), true)));

					// Check at the end:
					repeatBlock.add(pattern(repeatBlock) {
										if (${once})
											break;
										else
											return ${failResult};
									});
				}

				// Commit pending operations:
				for (x in repeatCommit)
					block.add(x);
				repeatCommit.clear();

				// Then restore position as appropriate.
				if (production.repType.repeatable)
					block.add(Continue(SrcPos(), block));
				else
					block.add(Break(SrcPos(), block));

				block = repeatPrev;
				// Added here to ensure that declaration of "me" is before the loop.
				block.add(repeatRoot);
			}

			// Store a variable. Act according to the repetition currently in affect.
			void storeVariable(Str name, LocalVar var) {
				if (repeatRoot.empty) {
					variables.put(name, var);
				} else if (production.repType == RepType:repZeroOne) {
					var type = wrapMaybe(var.result.asRef(false));
					Var v(block, type, SStr(name), Actuals());
					variables.put(name, v.var);
					block.add(v);

					LocalVarAccess expr(SrcPos(), var);
					LocalVarAccess target(SrcPos(), v.var);
					repeatCommit << Operator(block, target, assignOperator(SStr("="), 0), expr);
				} else {
					var type = wrapArray(var.result.asRef(false));
					Var v(block, type, SStr(name), Actuals());
					variables.put(name, v.var);
					block.add(v);

					LocalVarAccess expr(SrcPos(), var);
					LocalVarAccess to(SrcPos(), v.var);
					repeatCommit << namedExpr(block, SrcPos(), "push", to, Actuals(expr));
				}
			}

			// Invoke a function when appropriate (i.e. when we are done with the current "loop")
			void invokeLater(Expr expr) {
				if (repeatRoot.empty)
					block.add(expr);
				else
					repeatCommit << expr;
			}
		}

		// Create a particular rule.
		void create(Rule rule, Production[] productions) {
			// Some functions might be defined outside of there.
			unless (fn = get(rule) as BSTreeFn)
				return;
			// Check if we created it. There might be other BSTreeFn:s in the name tree.
			if (fn.parentLookup !is parent)
				return;

			// We are good to go!
			FnBody body(fn, scope);
			fn.body = body;

			// What to return on failure?
			Expr failResult = if (isMaybe(body.type)) {
				NullExpr(SrcPos());
			} else {
				BoolLiteral(SrcPos(), false);
			};

			// Store the position in a local variable.
			Var posStorage(body, SStr("@prevPos"), namedExpr(body, SStr("pos"), LocalVarAccess(SrcPos(), body.parameters[0])));
			body.add(posStorage);

			// Assignment operator:
			OpInfo assignInfo = assignOperator(SStr("="), 0);

			// Sort productions according to their priority:
			productions.sort((a, b) => a.priority > b.priority);

			// Look at first-sets for all productions in each "chunk". If we have a collision there,
			// ask the user to disambiguate the productions.
			Set<Regex> seen;
			Int lastPrio = 0; // Note: Does not matter if it clashes with the first element.
			for (p in productions) {
				if (p.priority != lastPrio)
					seen.clear();
				lastPrio = p.priority;

				Set<Regex> first = grammar.first(p);
				// TODO: Detect empty productions and put them last, or at least treat them as if
				// they are clashing with everything.
				for (x in first) {
					if (!seen.put(x))
						throwPrioError(x, productions, lastPrio);
				}

				// Create a clause here.
				createMatch(body, rule, p, failResult);
				// Restore parse state.
				var posMember = namedExpr(body, SStr("pos"), LocalVarAccess(SrcPos(), body.parameters[0]));
				body.add(Operator(body, posMember, assignInfo, LocalVarAccess(SrcPos(), posStorage.var)));
			}

			// If we get to this point, then the parse failed. Return null or false.
			body.add(failResult);

			// print("For ${rule.name}: ${body}");
		}

		// Create code that matches a particular production.
		private void createMatch(FnBody body, Rule rule, Production production, Expr failResult) {
			SrcPos pos = parent.pos;
			if (x = production.type)
				pos = x.pos;

			if (production.parent)
				throw GrammarError(pos, "Parent productions are not supported by the backtracking parser.");

			// Separate block to avoid variable clashes between productions (for variables created early).
			ExprBlock block(pos, body);
			body.add(block);

			// TODO: It might be relevant to examine the first-sets before attempting to parse in
			// order to avoid creating unnecessary objects?

			ProdState state(pos, block, production, LocalVarAccess(SrcPos(), body.parameters[0]), failResult);

			// Add parameters to the state object.
			for (i, x in rule.params) {
				state.variables.put(x.name, body.parameters[i + 1]);
			}

			// Go through the tokens:
			for (id, token in production.tokens) {
				if (token.raw)
					throw GrammarError(state.pos, "Capturing a raw parse tree is not supported.");

				if (id == production.repStart) {
					state.onRepStart();
				} else if (id == production.repEnd) {
					state.onRepEnd();
				}

				parseToken(state, token);
			}

			if (production.tokens.count == production.repEnd) {
				state.onRepEnd();
			}

			// If we get here, we succeeded. Return success!
			if (isMaybe(body.type)) {
				state.block.add(Return(pos, state.block, state.variable("me")));
			} else {
				// Create "me" now if it was not created already.
				state.createMeVar();

				state.block.add(Return(pos, state.block, BoolLiteral(SrcPos(), true)));
			}
		}

		// Parse a single token, and update the state as required.
		void parseToken(ProdState state, Token token) {
			if (parseBoundToken(state, token))
				return;

			// Not bound. Just evaluate it.
			Condition cond = createCondition(matchExpr(state, token));
			if (cond as WeakMaybeCast)
				cond.name(SStr("@tmp"));

			If check(state.block, cond);
			state.block.add(check);

			ExprBlock trueBlock(SrcPos(), check.successBlock);
			check.success(trueBlock);
			state.block = trueBlock;

			if (token as RegexToken) {
				unless (var = cond.result)
					throw InternalError("Failed to fetch the result.");

				// Need to update the position.
				LocalVarAccess newPos(SrcPos(), var);
				state.block.add(namedExpr(state.block, state.pos, "update", state.state, Actuals(newPos)));

				addError(state, token.regex, check);
			}
		}

		private Bool parseBoundToken(ProdState state, Token token) {
			unless (target = token.target)
				return false;

			if (token.invoke.empty & !token.bound)
				return false;

			Expr match = matchExpr(state, token);
			if (!isMaybe(match.result.type))
				throw GrammarError(state.pos, "Can not bind the result from a void rule to a variable.");

			WeakMaybeCast cast(match);
			SStr targetName(target.name);
			cast.name(targetName);

			If check(state.block, cast);
			state.block.add(check);

			ExprBlock trueBlock(SrcPos(), check.successBlock);
			check.success(trueBlock);
			state.block = trueBlock;

			unless (var = cast.result)
				throw InternalError("Failed to fetch the result.");

			if (token as RegexToken) {
				// For Regex tokens, we need to extract the substring and update 'pos':
				var access = LocalVarAccess(SrcPos(), var);
				var expr = pattern(state.block) {
					var substr = ${state.state}.input.cut(${state.state}.pos, ${access});
					${state.state}.update(${access});
					substr;
				};
				Var strVar(state.block, targetName, expr);
				state.block.add(strVar);
				var = strVar.var;

				addError(state, token.regex, check);
			}

			if (invoke = token.invoke) {
				// Invoke a function.
				// Note: The scope is technically incorrect here. It should be evaluated in
				// the context of the production rather than wherever the grammar was
				// declared.
				Actuals actuals;
				actuals.add(state.variable("me"));
				actuals.add(LocalVarAccess(SrcPos(), var));
				state.invokeLater(namedExpr(state.block, state.pos, invoke, actuals));
			} else if (token.bound) {
				state.storeVariable(target.name, var);
			}

			true;
		}

		// Add error message to condition.
		private void addError(ProdState state, Regex regex, If check) {
			StrLiteral msg(SrcPos(), "Expected: \"${regex}\"");
			check.fail(namedExpr(state.block, SrcPos(), "error", state.state, Actuals(msg)));
		}

		// Create a condition that matches the token.
		private Expr matchExpr(ProdState state, Token token) {
			if (token as RegexToken) {
				RegexValue val(token.regex);
				Actuals actuals;
				actuals.add(namedExpr(state.block, state.pos, "input", state.state));
				actuals.add(namedExpr(state.block, state.pos, "pos", state.state));
				return namedExpr(state.block, state.pos, "match", val, actuals);
			} else if (token as RuleToken) {
				Function toCall = get(token.rule);

				Actuals actuals;
				actuals.add(state.state);
				if (params = token.params) {
					for (param in params) {
						Expr src = state.variable(param);
						actuals.add(src);
					}
				}

				checkFnParams(token.rule, toCall, actuals, state.pos);

				return FnCall(SrcPos(), state.block.scope, toCall, actuals);
			} else {
				throw InternalError("Unknown token type: ${token}.");
			}
		}

		private void throwPrioError(Regex regex, Production[] productions, Int prio) {
			StrBuf msg;
			msg << "The grammar contains ambiguous productions with the same priority. ";
			msg << "While this works, it will not give predictable results. Therefore, ";
			msg << "you need to disambiguate this case by specifying priorities on the ";
			msg << "affected rules.\n";
			msg << "The following productions may all start with \"" << regex;
			msg << "\", and they all have the priority: " << prio << "\n";
			for (x in productions) {
				if (x.priority != prio)
					continue;
				if (grammar.first(x).has(regex))
					msg << x << "\n";
			}
			throw GrammarError(parent.pos, msg.toS);
		}

		// Load the implementation for the special productions.
		private void loadSpecial() {
			Type stateT = stateType(binary);

			// Note: We don't need to call 'loadAll' here. We don't need to load the special
			// functions if the grammar refers to them. If it does, it will be loaded at this
			// time. If the package is not loaded, we know that we don't need special functions, and
			// we thus don't need to force loading speicals.
			Package special = named{parser:special};
			for (named in special) {
				unless (named as Function)
					continue;
				unless (named.params.count > 0)
					continue;
				unless (named.params[0].mayStore(Value(stateT)))
					continue;

				// We have something that likely matches a rule now. Find the rule!
				unless (rule = special.find(SimplePart(named.name), Scope()) as Rule)
					continue;

				// Check so that parameters match.
				var params = rule.params;
				if (params.count + 1 != named.params.count)
					continue;
				Bool match = true;
				for (Nat i = 0; i < params.count; i++) {
					match &= params[i].type.type is named.params[i + 1].type;
				}
				if (!match)
					continue;

				if (resultType(rule.result).type !is named.result.type)
					continue;

				// Good match, add it!
				ruleFns.put(rule, named);
			}
		}
	}
}