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

/**
 * The recursive descent parser.
 *
 * It generates a set of functions, one for each rule, with the following structure:
 * T? Rx(State state, ...) {
 *     if (first-set of P1 matches) {
 *         // Update state according to regex matches or call any other rule fns.
 *         return me;
 *     }
 *     if (first-set of P2 matches) {
 *         // as above
 *     }
 *     state.error = "Expected ...";
 *     return null;
 * }
 *
 * Each of the "update state" chunks is the logic for matching a single production. It has the
 * following structure:
 *
 * var a = if (x = regex.match(state.input, state.pos)) {
 *     Str matched = state.input.substr(state.pos, x);
 *     state.pos = x;
 *     matched;
 * } else {
 *     state.error = "Expected ...";
 *     return null;
 * }
 *
 * or
 *
 * var a = if (x = Px(state, ...)) {
 *     x;
 * } else {
 *     return null;
 * }
 */
class RecursiveParser 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);
		// print(grammar.toS);

		{
			var recursive = grammar.leftRecursive;
			if (recursive.any) {
				StrBuf msg;
				msg << "Left recursive grammars are not supported by the recursive descent 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.error.any)
						result.error = createError(state.error, state.pos);
					return result;
				});

		body;
	}

	/**
	 * State for the parser, when parsing strings.
	 */
	class StrState extends RecursiveStrState {
		Str error;

		init(Str input, Str:Iter start) {
			init(input, start) {}
		}

		// Set error.
		void error(Str error) : override {
			this.error = error;
		}

		// Helper function from the grammar to match a set of regexes.
		// Does not update 'pos' on a match.
		Bool matches(Regex[] candidates) {
			for (x in candidates) {
				if (x.match(input, pos))
					return true;
			}
			false;
		}
	}

	/**
	 * State for the parser, when parsing binary buffers.
	 */
	class BinaryState extends RecursiveBinaryState {
		Str error;

		init(Buffer input, Nat start) {
			init(input, start) {}
		}

		// Set error.
		void error(Str error) : override {
			this.error = error;
		}

		// Helper function from the grammar to match a set of regexes.
		// Does not update 'pos' on a match.
		Bool matches(Regex[] candidates) {
			for (x in candidates) {
				if (x.match(input, pos))
					return true;
			}
			false;
		}
	}


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

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

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

			// Find any productions with empty first-sets.
			Production? empty;

			// Look at the first sets for each production.
			// If we have duplicates, that means the grammar is not LL(1).
			Set<Regex> seen;
			for (p in productions) {
				Set<Regex> first = grammar.first(p);
				for (x in first) {
					if (!seen.put(x))
						throwLLError(x, productions);
				}

				// Don't create a clause if the first-set is empty.
				if (first.empty)
					empty = p;
				else
					createClause(body, first, rule, p);
			}

			// See if there is an empty rule. If so, run that now!
			if (empty) {
				body.add(createParse(body, rule, empty));
			} else {
				// Emit an error and return null.
				LocalVarAccess state(SrcPos(), body.parameters[0]);
				StrBuf msg;
				msg << "Expected one of:";
				for (x in seen)
					msg << "\n\"" << x << "\"";
				StrLiteral msgStr(SrcPos(), msg.toS);
				body.add(pattern(body) {
							${state}.error = ${msgStr};
							return ${errorResult(body.type)};
						});
			}

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

		// Create a clause for a production inside the function.
		private void createClause(FnBody body, Set<Regex> first, Rule rule, Production production) {
			If check(body, createCondition(body, first));
			body.add(check);

			check.success(createParse(body, rule, production));
		}

		// Create code that evaluates to 'true' if one of the regexes match.
		private Expr createCondition(FnBody body, Set<Regex> first) {
			Regex[] candidates;
			for (x in first)
				candidates << x;

			LocalVarAccess stateVar(SrcPos(), body.parameters[0]);
			namedExpr(body, SrcPos(), "matches", stateVar, Actuals(RegexArray(candidates)));
		}

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

			// For error messages.
			SrcPos pos;

			// Root block to generate code into. Variables are always placed here, block may refer
			// to a sub-block from this (e.g. when in repeats).
			ExprBlock rootBlock;

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

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

			// Value to return on error.
			Expr errorVal;

			// Current production.
			Production production;

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

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

			// Are we inside the repetition? If so, this contains the expression we shall add with a
			// repetition in it.
			Expr? inRepeat;

			// Is the variable "me" 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;

				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(rootBlock, pos, result, actuals);
					if (init.result.type == Value()) {
						// It evaluates to 'void'. Make sure to stille valuate it.
						rootBlock.add(init);
						return false;
					}

					Var v(rootBlock, SStr("me"), init);
					rootBlock.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;
			}

			// Save a token to a variable, or invoke a function as appropriate.
			void storeToken(Token token, Expr expr) {
				if (target = token.target) {
					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(variable("me"));
						actuals.add(expr);
						block.add(namedExpr(block, pos, invoke, actuals));
					} else if (token.bound) {
						// Bind to a variable.
						storeVariable(target.name, expr);
					} else {
						// Evaluate the expression.
						block.add(expr);
					}
				}
			}

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

					block.add(pattern(block) { ${LocalVarAccess(SrcPos(), v.var)} = ${expr}; });
				} else {
					var type = wrapArray(expr.result.type.asRef(false));
					Var v(rootBlock, type, SStr(name), Actuals());
					variables.put(name, v.var);
					rootBlock.add(v);

					block.add(pattern(block) { ${LocalVarAccess(SrcPos(), v.var)}.push(${expr}); });
				}
			}

			// Figure out if we need to do anything at the start of a repetition.
			void onRepStart(Grammar g) {
				if (production.repCapture.any) {
					Var x(rootBlock, SStr("@capture start"), pattern(block) { ${state}.pos; });
					rootBlock.add(x);
					captureStart = x.var;
				} else if (production.repType == RepType:repZeroOne) {
					If c(rootBlock, createCond(g, production.repStart));
					inRepeat = c;

					block = ExprBlock(SrcPos(), c.successBlock);
					c.success(block);
				} else if (production.repType == RepType:repOnePlus) {
					Loop loop(SrcPos(), rootBlock);
					inRepeat = loop;
					loop.condExpr(createCond(g, production.repStart));

					block = ExprBlock(SrcPos(), loop);
					loop.doBody(block);
				} else if (production.repType == RepType:repZeroPlus) {
					Loop loop(SrcPos(), rootBlock);
					inRepeat = loop;
					loop.condExpr(createCond(g, production.repStart));

					block = ExprBlock(SrcPos(), loop);
					loop.whileBody(block);
				}
			}

			// Figure out if we need to do anything at the end of a repetition.
			void onRepEnd() {
				if (capture = production.repCapture) {
					unless (captureStart)
						throw InternalError("Captured portion of the production ${production} is inconsistent.");

					LocalVarAccess start(SrcPos(), captureStart);
					Expr captured = pattern(block) {
						${state}.input.cut(${start}, ${state}.pos);
					};
					storeToken(capture, captured);
				} else if (r = inRepeat) {
					rootBlock.add(r);
					inRepeat = null;
					block = rootBlock;
				}
			}

			// Create a condition that checks for the first-set of the token at position.
			private Expr createCond(Grammar grammar, Nat position) {
				Regex[] candidates;
				Token token = production.tokens[position];
				if (token as RuleToken) {
					for (x in grammar.first[token.rule])
						candidates << x;
				} else if (token as RegexToken) {
					candidates << token.regex;
				}

				namedExpr(block, SrcPos(), "matches", state, Actuals(RegexArray(candidates)));
			}
		}

		private Expr createParse(FnBody body, Rule rule, Production production) {
			// TODO: It might be good to only use the production's position if it is declared inline?
			SrcPos pos = parent.pos;
			if (type = production.type)
				pos = type.pos;

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

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

			// 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(grammar);
				} else if (id == production.repEnd) {
					state.onRepEnd();
				}

				if (target = token.target) {
					Expr e = parseToken(state, token, token.invoke.any | token.bound);
					state.storeToken(token, e);
				} else {
					state.block.add(parseToken(state, token, false));
				}
			}

			// End of the repeition might be after the last token.
			if (production.tokens.count == production.repEnd) {
				state.onRepEnd();
			}

			// print(production.toS);
			// print(block.toS);

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

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

			block;
		}

		// Create a suitable result value depending on the return type of the function.
		private Expr errorResult(Value returnType) {
			if (isMaybe(returnType))
				NullExpr(SrcPos());
			else
				BoolLiteral(SrcPos(), false);
		}

		// Parse a token:
		// - 'state' is the state variable
		// - 'needResult' is whether or not we need to capture the value from this token in the result.
		private Expr parseToken(ProdState state, Token token, Bool needResult) {
			if (token as RegexToken) {
				parseRegex(state, token.regex, needResult);
			} else if (token as RuleToken) {
				parseRule(state, token, needResult);
			} else {
				throw InternalError("Unknown token type: ${token}.");
			}
		}

		// "parse" a regex into some kind of value. Returns the transformed value.
		private Expr parseRegex(ProdState s, Regex regex, Bool needResult) {
			var state = s.state;
			StrLiteral msg(SrcPos(), "Expected something matching: ${regex}");
			if (needResult) {
				pattern(s.block) {
					if (x = ${RegexValue(regex)}.match(${state}.input, ${state}.pos)) {
						var substr = ${state}.input.cut(${state}.pos, x);
						${state}.pos = x;
						substr;
					} else {
						${state}.error = ${msg};
						return ${s.errorVal};
					}
				};
			} else {
				pattern(s.block) {
					if (x = ${RegexValue(regex)}.match(${state}.input, ${state}.pos)) {
						${state}.pos = x;
					} else {
						${state}.error = ${msg};
						return ${s.errorVal};
					}
				};
			}
		}

		// "parse" a rule.
		private Expr parseRule(ProdState state, RuleToken token, Bool needResult) {
			Function toCall = get(token.rule);

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

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

			FnCall fnCall(SrcPos(), state.block.scope, toCall, actuals);

			Condition cond = if (isMaybe(toCall.result)) {
				WeakMaybeCast x(fnCall);
				x.name(SStr("x"));
				x;
			} else {
				BoolCondition(fnCall);
			};
			If check(state.block, cond);
			if (x = cond.result) {
				check.success(LocalVarAccess(SrcPos(), x));
			} else if (needResult) {
				throw GrammarError(state.pos, "Unable to capture the value of the void-rule ${token.rule.identifier} in a production.");
			} else {
				// We don't need a result.
			}

			// False branch: just return the error message.
			check.fail(Return(SrcPos(), state.block, state.errorVal));

			check;
		}

		// Throw an error indicating problems with first-sets.
		// We don't worry about efficiency here, as we know we will fail compilation anyway.
		private void throwLLError(Regex regex, Production[] productions) {
			StrBuf msg;
			msg << "The grammar is not LL(1) since multiple productions may start with the same regex.\n";
			msg << "The following productions may all start with \"" << regex << "\":\n";
			for (x in productions) {
				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);
			}
		}
	}
}

// Helper to check function parameters. Assumes that the first parameter is not visible to the user.
package void checkFnParams(Rule rule, Function toCall, Actuals actuals, SrcPos pos) on Compiler {
	if (toCall.params.count != actuals.count) {
		StrBuf msg;
		msg << "Incorrect number of parameters to rule " << rule.name << ". ";
		msg << "Expected " << (toCall.params.count - 1) << " but got ";
		msg << (actuals.count - 1) << ".";
		throw GrammarError(pos, msg.toS);
	}

	for (i, actual in actuals.expressions) {
		if (!actual.castable(toCall.params[i], Scope())) {
			StrBuf msg;
			msg << "Type mismatch for parameter " << i << " in call to rule " << rule.name << ". ";
			msg << "Expected type " << toCall.params[i] << ", but got " << actual.result << ".";
			throw GrammarError(pos, msg.toS);
		}
	}
}