use core:lang;
use lang:bs;
use lang:bs:macro;

private Expr bindParam(SrcPos pos, Block block, Expr statement, Nat id, Expr param) on Compiler {
	Actuals params;
	params.add(NumLiteral(pos, id));
	params.add(param);
	namedExpr(block, pos, "bind", statement, params);
}

// Common parts of the code to create queries.
private Expr createQueryExpr(SrcPos pos, Block block, Expr connection, Query query, DatabaseType? type) on Compiler {
	if (type) {
		query.resolve(ResolveContext(block, type.contents));
	} else {
		query.resolve(ResolveContext(block));
	}

	var q = query.build();
	var cached = q.cachedQuery();

	if (type)
		type.queries << cached;

	ExprBlock r(pos, block);

	// Create a prepared statement.
	var prepared = {
		// Call connection.prepare(query). Note: It either refers to the one in
		// DBConnection or the one in TypedConnection, but they work the same way.
		Actuals params(lang:bs:util:ObjectLiteral(pos, cached));
		var stmt = namedExpr(r, pos, "prepare", connection, params);
		var varDecl = Var(r, SStr("statement"), stmt);
		r.add(varDecl);
		LocalVarAccess(pos, varDecl.var);
	};

	// Bind parameters.
	for (i, x in q.bind) {
		r.add(bindParam(pos, r, prepared, i, x));
	}

	// Execute it, and save the result.
	var resultDecl = Var(r, SStr("result"), namedExpr(r, pos, "execute", prepared));
	r.add(resultDecl);

	// Get the result, if desired. Otherwise, finalize it.
	LocalVarAccess resultAccess(pos, resultDecl.var);
	if (e = query.result(r, resultAccess))
		r.add(e);
	else
		r.add(namedExpr(r, pos, "finalize", resultAccess));

	r;
}


// Create a suitable query expression depending on what database connection was used.
Expr queryExpr(SrcPos pos, Block block, Expr expr, Query query) on Compiler {
	var result = expr.result().type();
	if (x = result.type as DatabaseType) {
		// Typed version.
		return createQueryExpr(pos, block, expr, query, x);
	} else if (Value(named{DBConnection}).mayReferTo(result)) {
		// Untyped version.
		return createQueryExpr(pos, block, expr, query, null);
	} else {
		throw SyntaxError(expr.pos, "Expected a database connection (typed or untyped).");
	}
}

// A Query block.
class QueryBlock extends ExprBlock {
	init(SrcPos pos, Block parent, Expr db) {
		var result = db.result.type;
		if (Value(named{DBConnection}).mayReferTo(result)) {
		} else if (result.type as DatabaseType) {
		} else {
			throw SyntaxError(db.pos, "Expected a database connection (typed or untyped).");
		}

		super(pos, parent);

		Var var(this, SStr(" db", db.pos), db);
		add(var);

		init {
			connection(db.pos, var.var);
		}
	}

	// Variable containing our database connection.
	LocalVarAccess connection;

	// Helper to find a Query block.
	QueryBlock find(Block in) : static {
		NameLookup? at = in.lookup;
		while (l = at as BlockLookup) {
			if (q = l.block as QueryBlock)
				return q;

			at = l.parent;
		}

		throw InternalError("Could not find a parent QueryBlock!");
	}
}

// Create a suitable query expression, assuming we're inside a WITH block.
Expr queryBlockExpr(SrcPos pos, Block block, Query query) on Compiler {
	QueryBlock q = QueryBlock:find(block);
	queryExpr(pos, block, q.connection, query);
}

/**
 * Table name. Optionally creates an alias for the query.
 */
class TableName on Compiler {
	SrcPos pos;
	Str name;
	Str alias;

	// No alias.
	init(SStr name) {
		init { pos = name.pos; name = name.v; alias = name.v; }
	}

	// Create an alias.
	init(SStr name, SStr alias) {
		init { pos = name.pos; name = name.v; alias = alias.v; }
	}

	// Build.
	void build(QueryBuilder b) {
		b.query.name(name);
		if (name != alias) {
			b.query.put(" AS ");
			b.query.name(alias);
		}
	}

	// To string.
	void toS(StrBuf to) {
		to << name;
		if (name != alias)
			to << " AS " << alias;
	}
}

/**
 * Context passed around while resolving a SQL query.
 */
class ResolveContext on Compiler {
	// Current BS block.
	Block block;

	// Contents of the database.
	Database? db;

	// Which tables are visible in the current context? Keys are aliases and not necessarily table names.
	Str->Table visible;

	// Which tables might be null due to a join? Keys are keys in 'visible'.
	Set<Str> nullTables;

	// Create typed version.
	init(Block block, Database db) {
		init { block = block; db = db; }
	}

	// Create untyped version.
	init(Block block) {
		init { block = block; }
	}

	// Is this a typed query?
	Bool typed() {
		db.any;
	}

	// Add a table to 'current'. Throws if the table is known not to exists.
	Table? addTable(SrcPos pos, Str to) {
		// Untyped version?
		unless (db)
			return null;

		unless (found = db.find(to))
			throw NoSuchTable(pos, to);
		if (visible.has(to))
			throw DuplicateAlias(pos, to);
		visible.put(to, found);
		found;
	}
	Table? addTable(TableName table) {
		unless (db)
			return null;

		unless (found = db.find(table.name))
			throw NoSuchTable(table.pos, table.name);
		if (visible.has(table.alias))
			throw DuplicateAlias(table.pos, table.alias);
		visible.put(table.alias, found);
		found;
	}

	// Result from "resolve"
	value ResolveResult {
		Str table;
		Column column;

		init(Str table, Column column) {
			init { table = table; column = column; }
		}
	}

	// Resolve a column.
	ResolveResult? resolve(SrcPos pos, Str? table, Str column) {
		if (table) {
			unless (t = visible.at(table))
				return null;
			if (col = t.find(column))
				return ResolveResult(table, col);
			return null;
		} else {
			ResolveResult? result;
			for (alias, t in visible) {
				if (col = t.find(column)) {
					if (result) {
						throw SyntaxError(pos, "The column name ${column} is ambigous. It exists in tables ${result.table} and ${t.name} at least.");
					} else {
						result = ResolveResult(alias, col);
					}
				}
			}
			return result;
		}
	}
}


/**
 * Query builder. Collects a SQL query string in a StrBuf, as well as BS expressions for any
 * parameters that need to be bound.
 */
class QueryBuilder on Compiler {
	// String builder that contains the final SQL query.
	QueryStrBuilder query;

	// Expressions that are used to bind parameters.
	Expr[] bind;

	// Append a bound parameter.
	void bind(Expr expr) {
		query.placeholder();
		bind << expr;
	}

	// Get a CachedQuery from the current state.
	CachedQuery cachedQuery() {
		SimpleCachedQuery(query.build);
	}

	// To string.
	void toS(StrBuf to) : override {
		to << "Query: " << query.build;
		to << "\nData: " << bind;
	}
}


/**
 * Query builder with support for fallbacks.
 */
class FallbackQueryBuilder extends QueryBuilder {
	// Features required to use the regular query.
	DBFeatures features;

	// Which query generates the result?
	Nat resultId;

	// Fallback queries.
	QueryStr[] fallbacks;

	// Create.
	init(DBFeatures features, Nat resultId) {
		init { features = features; resultId = resultId; }
	}

	// Cached query.
	CachedQuery cachedQuery() : override {
		MultiCachedQuery(query.build, features, fallbacks, resultId);
	}

	// To string.
	void toS(StrBuf to) : override {
		super:toS(to);
		to << "\nFeatures: " << features.toS;
		to << "\nFallbacks: " << join(fallbacks, ", ");
	}
}


/**
 * Base class for an SQL query.
 *
 * Note: .toS() looks like SQL, but does not take proper care of edge-cases, so should not be
 * treated as a proper SQL statement.
 */
class Query on Compiler {
	// Position in source.
	SrcPos pos;

	// Create.
	init(SrcPos pos) {
		init() {
			pos = pos;
		}
	}

	// Resolve any unknowns in this query, given a suitable context.
	// Also typechecks as applicable.
	void resolve(ResolveContext context) : abstract;

	// Build a query.
	QueryBuilder build() : abstract;

	// Compute the result of this query, if it should be available.
	Expr? result(Block context, Expr result) {
		null;
	}
}

/**
 * INSERT query.
 */
class InsertQuery extends Query {
	// Table name.
	SStr table;

	// Columns to insert into. If empty, we insert into all columns.
	SStr[] columns;

	// Values to insert.
	SQLExpr[] values;

	// Returning, if any.
	ReturningClause? returning;

	// Resolved table, after 'resolve' is called.
	private Table? resolvedTable;

	// Create, insert all columns.
	init(SrcPos pos, SStr table, SQLExpr[] values, ReturningClause? returning) {
		init(pos) {
			table = table;
			values = values;
			returning = returning;
		}
	}

	// Create, insert into only a subset of columns.
	init(SrcPos pos, SStr table, SStr[] columns, SQLExpr[] values, ReturningClause? returning) {
		init(pos) {
			table = table;
			columns = columns;
			values = values;
			returning = returning;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		var table = context.addTable(table.pos, table.v);
		this.resolvedTable = table;

		// If we found a table, explicitly specify the column names for easier typechecking later.
		Column[] cols;
		if (table) {
			if (columns.empty) {
				for (c in table.columns)
					columns << SStr(c.name, pos);
			} else {
				// Find all of the columns.
				Set<Str> used;
				for (c in columns) {
					used.put(c.v);
					unless (found = table.find(c.v))
						throw NoSuchColumn(c.pos, c.v, table.name);
					cols << found;
				}

				Bool multiPK = table.multiplePK();

				// Check the ones that are not inserted.
				for (c in table.columns) {
					if (used.has(c.name))
						continue;

					if (!c.hasDefault(multiPK))
						throw SyntaxError(pos, "The column ${c.name} has no default value, and needs to be given a value.");
				}
			}

			if (returning) {
				returning.resolve(context, table);

				// We need to reorder the columns of the table to match the order they are selected.
				// Otherwise, any bound parameters will not line up properly.
				Column[] orderedCols;
				SQLExpr[] orderedVals;
				for (desc in returning.rowDesc) {
					for (i, x in cols) {
						if (x.name == desc.name) {
							orderedCols << x;
							orderedVals << values[i];
							cols.remove(i);
							values.remove(i);
							break;
						}
					}
				}

				// Copy remaining ones.
				for (i, x in cols) {
					orderedCols << x;
					orderedVals << values[i];
				}

				cols = orderedCols;
				values = orderedVals;

				// Copy names to 'columns' since we changed the order.
				columns.clear();
				for (x in cols)
					columns << SStr(x.name);
			}
		}

		if (cols.any & (values.count != cols.count))
			throw SyntaxError(pos, "The number of values does not match the number of columns inserted into.");

		// It does not really make sense to refer to column names in an insert statement...
		context.visible.clear();

		for (Nat i = 0; i < values.count; i++)
			values[i] = values[i].resolve(context);

		for (i, col in cols) {
			Value result(values[i].result);
			if (col.allowNull)
				result = unwrapMaybe(result);

			if (!Value(col.datatype.storm).mayReferTo(result))
				throw SyntaxError(values[i].pos, "Can not store a value of type ${result} in the column \"${col.name}\".");
		}
	}

	// Build the query.
	QueryBuilder build() : override {
		if (returning) {
			// Note: We can not support RETURNING for non-typed databases!
			if (returning.rowDesc.empty)
				throw SyntaxError(returning.pos, "RETURNING statement can only be used for typed connections.");

			FallbackQueryBuilder r(DBFeatures:insertReturning, 1);
			buildNative(r);

			QueryStr part1 = r.query.build();

			r.query.put(" ");
			returning.build(r);

			QueryBuilder part2;
			part2.query.put("SELECT ");
			for (i, col in returning.rowDesc) {
				if (i > 0)
					part2.query.put(", ");
				putColValue(part2, col.name);
			}

			r.fallbacks << part1;
			r.fallbacks << part2.query.build();

			r;
		} else {
			QueryBuilder r;
			buildNative(r);
			r;
		}
	}

	// Find the value of a column.
	private void putColValue(QueryBuilder r, Str column) {
		unless (resolvedTable)
			return;

		// See if this is a column we inserted into:
		if (columns.any) {
			for (i, c in columns) {
				if (c.v == column) {
					values[i].build(r);
					return;
				}
			}
		}

		// Find the column:
		for (i, col in resolvedTable.columns) {
			if (col.name != column)
				continue;

			// If we insert into all columns, it is not caught from above. Check here instead.
			if (columns.empty) {
				values[i].build(r);
				return;
			}

			// Figure out the default value:
			if (d = col.default) {
				d.build(r);
				return;
			}

			// Is it an autoincrement column?
			Bool autoIncrement = col.autoIncrement;
			if (!autoIncrement)
				autoIncrement = i == resolvedTable.implicitAutoIncrementColumn();
			if (autoIncrement) {
				r.query.lastRowId();
				return;
			}

			// Nothing else, just put null.
			r.query.put("NULL");
			return;
		}
	}

	// Create the "native" query.
	private void buildNative(QueryBuilder r) {
		r.query.put("INSERT INTO ");
		r.query.name(table.v);
		if (columns.any) {
			r.query.put(" (");
			for (i, col in columns) {
				if (i > 0)
					r.query.put(", ");
				r.query.name(col.v);
			}
			r.query.put(")");
		}
		r.query.put(" VALUES (");
		for (i, v in values) {
			if (i > 0)
				r.query.put(", ");
			v.build(r);
		}
		r.query.put(")");
	}

	// Return the last created row ID.
	Expr? result(Block context, Expr result) {
		if (returning) {
			returning.singleResult(context, result);
		} else {
			namedExpr(context, pos, "lastRowId", result);
		}
	}

	// To string.
	void toS(StrBuf to) : override {
		to << "INSERT INTO " << table.v;
		if (columns.any)
			to << " (" << join(columns, ", ", (x) => x.v) << ")";

		to << " VALUES (" << join(values, ", ") << ");";
	}
}

// Helper to check that a condition returns a boolean. Assumes that the condition was resolved beforehand.
private void checkCondition(SQLExpr condition) on Compiler {
	var result = condition.result;
	if (!Value(named{Bool}).mayReferTo(unwrapMaybe(result)))
		throw SyntaxError(condition.pos, "Expressions in WHERE and ON clauses are expected to return a Bool, not ${result}.");
}

/**
 * UPDATE query.
 */
class UpdateQuery extends Query {
	// Table to update.
	SStr table;

	// Columns to update.
	AssignExpr[] update;

	// Condition, if any.
	SQLExpr? condition;

	// Returning clause, if any.
	ReturningClause? returning;

	// Create.
	init(SrcPos pos, SStr table, AssignExpr[] update, SQLExpr? condition, ReturningClause? returning) {
		init(pos) {
			table = table;
			update = update;
			condition = condition;
			returning = returning;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		var table = context.addTable(table.pos, table.v);

		// Resolve all assignments.
		for (x in update) {
			x.value = x.value.resolve(context);
		}

		// Update the expression if suitable.
		if (x = condition) {
			var resolved = x.resolve(context);
			if (context.typed)
				checkCondition(resolved);
			condition = resolved;
		}

		if (table) {
			// Check if all columns exist, and type-check the assignments.
			for (x in update) {
				unless (column = table.find(x.column.v))
					throw NoSuchColumn(x.column.pos, x.column.v, table.name);

				Value result(x.value.result);
				if (!Value(column.datatype.storm).mayReferTo(result))
					throw SyntaxError(x.value.pos, "Cannot assign a ${result} to column ${column.name}");
			}
		}

		if (returning) {
			returning.resolve(context, table);

			// For fallbacks to work, we need to re-order the updates to be in the same order as the
			// returned columns. Otherwise, any bound parameters may not line up in the separate
			// fallback queries!
			AssignExpr[] reordered;
			for (x in returning.rowDesc) {
				for (i, y in update) {
					if (y.column.v == x.name) {
						reordered << y;
						update.remove(i);
						break;
					}
				}
			}
			update = reordered;
		}
	}

	// Build the query.
	QueryBuilder build() : override {
		if (returning) {
			// Note: We can not support RETURNING for non-typed databases!
			if (returning.rowDesc.empty)
				throw SyntaxError(returning.pos, "RETURNING statement can only be used for typed connections.");

			FallbackQueryBuilder r(DBFeatures:updateReturning, 0);
			buildNative(r);

			// Save the first part of the good query.
			QueryStr part2 = r.query.build();

			r.query.put(" ");
			returning.build(r);

			// Create a select query that selects updated values for columns.
			QueryBuilder part1;
			part1.query.put("SELECT ");
			for (id, column in returning.rowDesc) {
				if (id > 0)
					part1.query.put(", ");

				if (!putReplacement(part1, column))
					part1.query.name(column.name);
			}
			part1.query.put(" FROM ");
			part1.query.name(table.v);
			if (condition) {
				part1.query.put(" WHERE ");
				condition.build(part1);
			}

			r.fallbacks << part1.query.build();
			r.fallbacks << part2;

			r;
		} else {
			QueryBuilder r;
			buildNative(r);
			r;
		}
	}

	// Find a replacement expression for a column.
	private Bool putReplacement(QueryBuilder r, TypedCol column) {
		for (u in update) {
			if (u.column.v == column.name) {
				u.value.build(r);
				return true;
			}
		}
		return false;
	}

	// Create the "native" query.
	private void buildNative(QueryBuilder r) {
		r.query.put("UPDATE ");
		r.query.name(table.v);
		r.query.put(" SET ");
		for (i, x in update) {
			if (i > 0)
				r.query.put(", ");
			r.query.put(x.column.v);
			r.query.put(" = ");
			x.value.build(r);
		}

		if (condition) {
			r.query.put(" WHERE ");
			condition.build(r);
		}
	}

	// Return the number of modified rows.
	Expr? result(Block context, Expr result) {
		if (returning) {
			returning.result(context, result);
		} else {
			namedExpr(context, pos, "changes", result);
		}
	}

	// To string.
	void toS(StrBuf to) : override {
		to << "UPDATE " << table.v << " SET " << join(update, ", ");
		if (condition)
			to << " WHERE " << condition;
		if (returning)
			to << " " << returning;
	}
}

/**
 * Value to update in an UPDATE query.
 */
class AssignExpr on Compiler {
	// Column to update.
	SStr column;

	// Value to assign.
	SQLExpr value;

	// Create.
	init(SStr column, SQLExpr value) {
		init { column = column; value = value; }
	}

	// To string.
	void toS(StrBuf to) : override {
		to << column.v << " = " << value;
	}
}


/**
 * DELETE query.
 */
class DeleteQuery extends Query {
	// Table to update.
	SStr table;

	// Condition, if any.
	SQLExpr? condition;

	// Returning clause, if any.
	ReturningClause? returning;

	// Create.
	init(SrcPos pos, SStr table, SQLExpr? condition, ReturningClause? returning) {
		init(pos) {
			table = table;
			condition = condition;
			returning = returning;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		var table = context.addTable(table.pos, table.v);

		if (x = condition) {
			var resolved = x.resolve(context);
			if (context.typed)
				checkCondition(resolved);
			condition = resolved;
		}
		if (returning) {
			returning.resolve(context, table);
		}
	}

	// Build the query.
	QueryBuilder build() : override {
		if (returning) {
			FallbackQueryBuilder r(DBFeatures:deleteReturning, 0);
			buildNative(r);

			// Save the DELETE FROM query.
			QueryStr part2 = r.query.build();

			r.query.put(" ");
			returning.build(r);
			r.query.put(";");

			// Create a select query.
			QueryBuilder part1;
			part1.query.put("SELECT ");
			returning.build(part1.query);
			part1.query.put(" FROM ");
			part1.query.name(table.v);
			if (condition) {
				part1.query.put(" WHERE ");
				condition.build(part1);
			}

			r.fallbacks << part1.query.build();
			r.fallbacks << part2;

			r;
		} else {
			QueryBuilder r;
			buildNative(r);
			r;
		}
	}

	// Crate the "native" query.
	private void buildNative(QueryBuilder r) {
		r.query.put("DELETE FROM ");
		r.query.name(table.v);
		if (condition) {
			r.query.put(" WHERE ");
			condition.build(r);
		}
	}

	// Return the number of modified rows.
	Expr? result(Block context, Expr result) {
		if (returning) {
			returning.result(context, result);
		} else {
			namedExpr(context, pos, "changes", result);
		}
	}

	// To string.
	void toS(StrBuf to) : override {
		to << "DELETE FROM " << table.v;
		if (condition)
			to << " WHERE " << condition;
		if (returning)
			to << " " << returning;
	}
}


/**
 * Common parts for SELECT-like queries (SELECT, SELECT ONE and COUNT currently).
 *
 * Handles the logic of JOIN and WHERE clauses.
 */
class SelectBase extends Query {
	// Table.
	TableName table;

	/**
	 * Join statement.
	 */
	class Join on Compiler {
		// Type of join.
		JoinType type;

		// Table.
		TableName table;

		// Condition for join.
		SQLExpr expr;

		// Create.
		init(JoinType type, TableName table, SQLExpr expr) {
			init {
				type = type;
				table = table;
				expr = expr;
			}
		}

		// Build.
		void build(QueryBuilder to) {
			to.query.put(typeStr());
			to.query.put(" ");
			table.build(to);
			to.query.put(" ON ");
			expr.build(to);
		}

		// String of join type.
		Str typeStr() {
			if (type == JoinType:inner) {
				"JOIN";
			} else if (type == JoinType:full) {
				"FULL JOIN";
			} else if (type == JoinType:left) {
				"LEFT JOIN";
			} else if (type == JoinType:right) {
				"RIGHT JOIN";
			} else {
				"";
			}
		}

		// ToS.
		void toS(StrBuf to) : override {
			to << typeStr() << " " << table << " ON " << expr;
		}
	}

	// Type of joins.
	// Masked to easily mask out which tables might be null.
	enum JoinType : bitmask {
		inner = 0x0,
		full = 0x3,
		left = 0x2,
		right = 0x1,

		nullBefore = 0x1,
		nullAfter = 0x2,
	}

	// Join clauses.
	Join[] joins;

	// Where clause, if present.
	SQLExpr? condition;

	init(SrcPos pos, TableName table, Join[] joins, SQLExpr? condition) {
		init(pos) {
			table = table;
			joins = joins;
			condition = condition;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		context.addTable(table);

		for (j in joins) {
			if (j.type.has(JoinType:nullBefore)) {
				// For left and full, all tables before might be null.
				for (k, v in context.visible)
					context.nullTables.put(k);
			}

			context.addTable(j.table);
			if (j.type.has(JoinType:nullAfter)) {
				// For right and full, current table might be null.
				context.nullTables.put(j.table.alias);
			}

			// TODO: Might be too early, but I think this is how SQL works.
			j.expr = j.expr.resolve(context);
			if (context.typed)
				checkCondition(j.expr);
		}

		if (x = condition) {
			var resolved = x.resolve(context);
			if (context.typed)
				checkCondition(resolved);
			condition = resolved;
		}
	}

	// Build the tail of the query (i.e. JOIN and WHERE clauses).
	protected void buildTail(QueryBuilder to) {
		for (j in joins) {
			to.query.put(" ");
			j.build(to);
		}

		if (condition) {
			to.query.put(" WHERE ");
			condition.build(to);
		}
	}

	// To string. Once again, only the tail parts.
	protected void tailToS(StrBuf to) {
		for (j in joins)
			to << j;
		if (condition)
			to << " WHERE " << condition;
	}
}


/**
 * Column to extract.
 */
class SelectedColumn on Compiler {
	// Position
	SrcPos pos;

	// Table (optional). Note, we use the alias of the table if specified.
	Str? table;

	// Column.
	Str column;

	// Alternate name as used in the output.
	Str? as;

	// Type of this column, if available.
	Type? type;

	init(SrcPos pos, SStr column) {
		init {
			pos = pos;
			column = column.v;
		}
	}

	init(SrcPos pos, SStr table, SStr column) {
		init {
			pos = pos;
			table = table.v;
			column = column.v;
		}
	}

	init(SrcPos pos, Str? table, Str column) {
		init {
			pos = pos;
			table = table;
			column = column;
		}
	}

	// Set alternate name in the output.
	void setAs(SStr to) {
		as = to.v;
	}

	// Build.
	void build(QueryBuilder to) {
		build(to.query);
	}

	void build(QueryStrBuilder to) {
		if (table) {
			to.name(table);
			to.put(".");
		}
		to.name(column);
	}

	// ToS.
	void toS(StrBuf to) : override {
		if (table)
			to << table << ".";
		to << column << " AS " << as;
	}
}


/**
 * Returning clause used for INSERT, DELETE, and UPDATE.
 */
class ReturningClause on Compiler {
	// Position.
	SrcPos pos;

	// Selected columns.
	SelectedColumn[] columns;

	// Description of the types in each row. Empty for non-typed queries.
	TypedCol[] rowDesc;

	// Create.
	init(SrcPos pos, SelectedColumn[] cols) {
		init { columns = cols; }
	}

	// Resolve columns.
	void resolve(ResolveContext context, Table? mainTable) {
		if (context.visible.any) {
			rowDesc.clear();
			if (columns.empty) {
				addAllCols(context);
			} else {
				resolveCols(context, mainTable);
			}
		}
	}

	// Add all columns to our output.
	private void addAllCols(ResolveContext context) {
		for (name, table in context.visible) {
			NestedTypedCol nested(name, false);

			for (id, col in table.columns) {
				SelectedColumn c(pos, null, col.name);
				Type t = resultType(false, col);
				c.type = t;
				nested.push(ScalarTypedCol(col.name, id, t));
			}

			rowDesc << nested;
		}

		// Remove scoping if we only had one table:
		if (rowDesc.count == 1) {
			if (nested = rowDesc[0] as NestedTypedCol) {
				rowDesc = nested.columns;
			}
		}
	}

	// Resolve columns and create our rowDesc.
	private void resolveCols(ResolveContext context, Table? table) {
		Nat id = 0;
		Str->NestedTypedCol nested;

		for (col in columns) {
			unless (result = context.resolve(col.pos, col.table, col.column)) {
				if (table = col.table)
					throw NoSuchColumn(col.pos, col.column, table);
				else if (table)
					throw NoSuchColumn(col.pos, col.column, table.name);
				else
					throw NoSuchColumn(col.pos, col.column, "<unknown table>");
			}

			Bool nullTable = context.nullTables.has(result.table);

			// Now, find the type and add it to the rowDesc. We need to wrap the type inside Maybe
			// only when we put it in a context where it is not scoped.
			if (altName = col.as) {
				// If it has an alternate name, it is never scoped.
				Type t = resultType(nullTable, result.column);
				col.type = t;
				rowDesc << ScalarTypedCol(altName, id++, t);
			} else if (table = col.table) {
				// A table was specified originally. Always scoped, no need for extra Maybe.
				Type t = resultType(false, result.column);
				col.type = t;

				if (!nested.has(table)) {
					NestedTypedCol created(table, nullTable);
					nested.put(table, created);
					rowDesc << created;
				}
				nested.get(table).push(ScalarTypedCol(col.column, id++, t));
			} else {
				// Else, the table was not specified originally and the result will not be scoped.
				Type t = resultType(nullTable, result.column);
				col.type = t;
				rowDesc << ScalarTypedCol(col.column, id++, t);
			}
		}
	}

	// Figure out the type for a column.
	private Type resultType(Bool nullTable, Column col) {
		Type t = col.datatype.storm;
		if (nullTable | col.allowNull) {
			if (maybe = wrapMaybe(t).type)
				return maybe;
		}
		t;
	}

	// Generate result.
	Expr result(Block context, Expr result) {
		if (rowDesc.empty)
			return namedExpr(context, pos, "iter", result);

		var iterType = getTypedIter(rowDesc);
		unless (ctor = iterType.find("__init", [Value(iterType), result.result.type.asRef(false)], Scope()) as Function)
			throw InternalError("Could not find a suitable constructor in the generated type.");

		CtorCall(pos, context.scope, ctor, Actuals(result));
	}

	// Generate a single result.
	Expr singleResult(Block context, Expr result) {
		// No type info - just return a Maybe<Row> object.
		if (rowDesc.empty)
			return namedExpr(context, pos, "next", result);

		// Otherwise, we want to return either <null> or an instance of the row-type.
		TypedRow t = getTypedRow(rowDesc);
		unless (ctor = t.find("__init", [Value(t), Value(named{Row})], Scope()) as Function)
			throw InternalError("Could not find a suitable constructor in the generated type.");

		unless (maybe = wrapMaybe(Value(t)).type)
			throw InternalError("Could not find the maybe type for the generated row type.");

		// Create an if-statement.
		WeakMaybeCast cast(namedExpr(context, pos, "next", result));
		cast.name(SStr("x"));
		If check(context, cast);

		if (created = cast.result) {
			check.success(CtorCall(pos, Scope(), ctor, Actuals(LocalVarAccess(pos, created))));
		}

		if (maybeCtor = maybe.find("__init", [Value(maybe, true)], Scope()) as Function) {
			check.fail(CtorCall(pos, Scope(), maybeCtor, Actuals()));
		}

		check;
	}

	// Build a query.
	void build(QueryBuilder out) {
		out.query.put("RETURNING ");
		build(out.query);
	}

	void build(QueryStrBuilder out) {
		if (columns.empty) {
			out.put("*");
		} else {
			for (i, col in columns) {
				if (i > 0)
					out.put(", ");
				col.build(out);
			}
		}
	}

	// ToS.
	void toS(StrBuf to) : override {
		to << "RETURNING ";
		if (columns.empty)
			to << "*";
		else
			to << join(columns, ", ");
	}
}


/**
 * SELECT query.
 */
class SelectQuery extends SelectBase {
	/**
	 * Single order by column.
	 */
	class OrderBy on Compiler {
		// Position
		SrcPos pos;

		// Table (optional).
		Str? table;

		// Column.
		Str column;

		// Ascending/descending.
		Bool ascending;

		init(SrcPos pos, SStr? table, SStr column, Bool ascending) {
			Str? t;
			if (table)
				t = table.v;

			init {
				pos = pos;
				table = t;
				column = column.v;
				ascending = ascending;
			}
		}

		// Build.
		void build(QueryBuilder to) {
			if (table) {
				to.query.name(table);
				to.query.put(".");
			}
			to.query.name(column);

			if (ascending)
				to.query.put(" ASC");
			else
				to.query.put(" DESC");
		}

		// To string.
		void toS(StrBuf to) {
			if (table)
				to << table << ".";
			to << column;
			if (ascending)
				to << " ASC";
			else
				to << " DESC";
		}
	}

	// Columns to select. If empty, we assume all columns.
	SelectedColumn[] cols;

	// Order by.
	OrderBy[] orderBy;

	// Description of types in each row. Empty for non-typed queries.
	TypedCol[] rowDesc;

	// Create.
	init(SrcPos pos, TableName table, SelectedColumn[] cols, Join[] joins, SQLExpr? condition, OrderBy[] orderBy) {
		init(pos, table, joins, condition) {
			cols = cols;
			orderBy = orderBy;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		super:resolve(context);

		// Type-check and prepare output format:
		if (context.visible.any) {
			rowDesc.clear();
			if (cols.empty) {
				addAllCols(context);
			} else {
				resolveCols(context);
			}

			resolveOrder(context);
		}
	}

	// Add all columns to our output.
	private void addAllCols(ResolveContext context) {
		if (joins.empty) {
			// Only one table, don't scope names:
			var name = this.table.alias;
			var table = context.visible.get(name);

			// Note: Currently, nullTable should always be true. This is, however, here for
			// completeness if this assumption would change in the future.
			Bool nullTable = context.nullTables.has(name);

			for (id, col in table.columns) {
				SelectedColumn c(pos, null, col.name);
				Type t = resultType(nullTable, col);

				c.type = t;
				cols << c;

				rowDesc << ScalarTypedCol(col.name, id, t);
			}
		} else {
			// Mutliple tables, scope them appropriately.

			// Note: We could loop through tables in 'context', but to get the order consistent in
			// the output of row types we traverse 'table' and 'join' instead.
			rowDesc << addAllCols(this.table.alias, context);

			for (x in joins)
				rowDesc << addAllCols(x.table.alias, context);
		}
	}

	// Add all names from a table to our output.
	private NestedTypedCol addAllCols(Str tableAlias, ResolveContext context) {
		NestedTypedCol nested(tableAlias, context.nullTables.has(tableAlias));

		var table = context.visible.get(tableAlias);
		for (col in table.columns) {
			SelectedColumn c(pos, tableAlias, col.name);
			Type t = resultType(false, col);

			nested.push(ScalarTypedCol(col.name, cols.count, t));

			c.type = t;
			cols << c;
		}

		nested;
	}

	// Resolve columns and create our rowDesc.
	private void resolveCols(ResolveContext context) {
		Nat id = 0;
		Str->NestedTypedCol nested;

		for (col in cols) {
			unless (result = context.resolve(col.pos, col.table, col.column)) {
				if (table = col.table)
					throw NoSuchColumn(col.pos, col.column, table);
				else
					throw NoSuchColumn(col.pos, col.column, table.name);
			}

			Bool nullTable = context.nullTables.has(result.table);

			// Now, find the type and add it to the rowDesc. We need to wrap the type inside Maybe
			// only when we put it in a context where it is not scoped.
			if (altName = col.as) {
				// If it has an alternate name, it is never scoped.
				Type t = resultType(nullTable, result.column);
				col.type = t;
				rowDesc << ScalarTypedCol(altName, id++, t);
			} else if (table = col.table) {
				// A table was specified originally. Always scoped, no need for extra Maybe.
				Type t = resultType(false, result.column);
				col.type = t;

				if (!nested.has(table)) {
					NestedTypedCol created(table, nullTable);
					nested.put(table, created);
					rowDesc << created;
				}
				nested.get(table).push(ScalarTypedCol(col.column, id++, t));
			} else {
				// Else, the table was not specified originally and the result will not be scoped.
				Type t = resultType(nullTable, result.column);
				col.type = t;
				rowDesc << ScalarTypedCol(col.column, id++, t);
			}
		}
	}

	// Resolve the "order by" portion.
	private void resolveOrder(ResolveContext context) {
		for (col in orderBy) {
			if (context.resolve(col.pos, col.table, col.column).empty) {
				if (table = col.table)
					throw NoSuchColumn(col.pos, col.column, table);
				else
					throw NoSuchColumn(col.pos, col.column, table.name);
			}
		}
	}

	// Figure out the type for a column.
	private Type resultType(Bool nullTable, Column col) {
		Type t = col.datatype.storm;
		if (nullTable | col.allowNull) {
			if (maybe = wrapMaybe(t).type)
				return maybe;
		}
		t;
	}

	// Build the query.
	QueryBuilder build() : override {
		QueryBuilder r;

		r.query.put("SELECT ");
		if (cols.empty) {
			r.query.put("*");
		} else {
			for (i, name in cols) {
				if (i > 0)
					r.query.put(", ");
				name.build(r);
			}
		}
		r.query.put(" FROM ");
		table.build(r);

		buildTail(r);

		if (orderBy.any) {
			r.query.put(" ORDER BY ");
			for (i, name in orderBy) {
				if (i > 0)
					r.query.put(", ");
				name.build(r);
			}
		}

		r.query.put(";");
		r;
	}

	// Return a proper iterator.
	Expr? result(Block context, Expr result) {
		// If no type information, fall back to the typeless version.
		if (rowDesc.empty)
			return namedExpr(context, pos, "iter", result);

		var iterType = getTypedIter(rowDesc);
		unless (ctor = iterType.find("__init", [Value(iterType), result.result.type.asRef(false)], Scope()) as Function)
			throw InternalError("Could not find a suitable constructor in the generated type.");

		CtorCall(pos, context.scope, ctor, Actuals(result));
	}

	// To string.
	void toS(StrBuf to) : override {
		to << "SELECT " << join(cols, ", ") << " FROM ";
		table.toS(to);
		tailToS(to);
		if (orderBy.any) {
			to << " ORDER BY ";
			for (i, name in orderBy) {
				if (i > 0)
					to << ", ";
				name.toS(to);
			}
		}
	}
}

/**
 * SELECT ONE query.
 *
 * Works just like SELECT, but returns Maybe<Row> rather than an iterator.
 */
class SelectOneQuery extends SelectQuery {
	init(SrcPos pos, TableName table, SelectedColumn[] cols, Join[] joins, SQLExpr? condition, OrderBy[] orderBy) {
		init(pos, table, cols, joins, condition, orderBy) {}
	}

	// Modified result. Returns a single row (wrapped in Maybe) rather than an iterator.
	Expr? result(Block context, Expr result) {
		// No type info - just return a Maybe<Row> object.
		if (rowDesc.empty)
			return namedExpr(context, pos, "next", result);

		// Otherwise, we want to return either <null> or an instance of the row-type.
		TypedRow t = getTypedRow(rowDesc);
		unless (ctor = t.find("__init", [Value(t), Value(named{Row})], Scope()) as Function)
			throw InternalError("Could not find a suitable constructor in the generated type.");

		unless (maybe = wrapMaybe(Value(t)).type)
			throw InternalError("Could not find the maybe type for the generated row type.");

		// Create an if-statement.
		WeakMaybeCast cast(namedExpr(context, pos, "next", result));
		cast.name(SStr("x"));
		If check(context, cast);

		if (created = cast.result) {
			check.success(CtorCall(pos, Scope(), ctor, Actuals(LocalVarAccess(pos, created))));
		}

		if (maybeCtor = maybe.find("__init", [Value(maybe, true)], Scope()) as Function) {
			check.fail(CtorCall(pos, Scope(), maybeCtor, Actuals()));
		}

		check;
	}
}

/**
 * COUNT FROM query.
 *
 * Works like SELECT COUNT(*) FROM, but clearer. Returns a Nat.
 */
class CountQuery extends SelectBase {
	init(SrcPos pos, TableName table, Join[] joins, SQLExpr? condition) {
		init(pos, table, joins, condition) {}
	}

	// Build the query.
	QueryBuilder build() : override {
		QueryBuilder r;

		r.query.put("SELECT COUNT(*) FROM ");
		table.build(r);

		buildTail(r);

		r.query.put(";");
		r;
	}

	// Return a Nat representing the number of rows.
	Expr? result(Block context, Expr result) {
		pattern(context) {
			if (row = ${result}.next()) {
				row.getLong(0).nat;
			} else {
				// This should not happen...
				0;
			}
		};
	}

	// To string.
	void toS(StrBuf to) : override {
		to << "COUNT FROM ";
		table.toS(to);
		tailToS(to);
	}
}

/**
 * CREATE TABLE query.
 *
 * Only for untyped connections.
 */
class CreateQuery extends Query {
	// Table declaration.
	Table table;

	// Is this a "CREATE IF NOT EXISTS"?
	Bool ifNotExists;

	// Create.
	init(SrcPos pos, Bool ifNotExists, Table table) {
		init(pos) {
			table = table;
			ifNotExists = ifNotExists;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		if (context.typed)
			throw SyntaxError(pos, "Can not use CREATE TABLE queries with typed connections.");
	}

	// Build the query.
	QueryBuilder build() : override {
		QueryBuilder r;
		// Note: We give 'true' for 'supports automatic autoincrement' to avoid modifying the query.
		table.toSQL(r.query, ifNotExists, true, []);
		r;
	}

}

/**
 * CREATE INDEX query.
 *
 * Only for untyped connections.
 */
class IndexQuery extends Query {
	// Index name.
	Str name;

	// Table name.
	SStr table;

	// Columns.
	SStr[] columns;

	// Create.
	init(SrcPos pos, SStr name, SStr table, SStr[] columns) {
		init(pos) {
			name = name.v;
			table = table;
			columns = columns;
		}
	}

	// Create.
	init(SrcPos pos, SStr table, SStr[] columns) {
		StrBuf name;
		name << table.v << "_" << join(columns, "_", (x) => x.v);

		init(pos) {
			name = name.toS;
			table = table;
			columns = columns;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		if (context.typed)
			throw SyntaxError(pos, "Can not use CREATE TABLE queries with typed connections.");
	}

	// Build the query.
	QueryBuilder build() : override {
		QueryBuilder r;

		r.query.put("CREATE INDEX ");
		r.query.name(name);
		r.query.put(" ON ");
		r.query.name(table.v);
		r.query.put("(");
		for (i, c in columns) {
			if (i > 0)
				r.query.put(", ");
			r.query.name(c.v);
		}
		r.query.put(");");

		r;
	}

}

/**
 * DROP TABLE query.
 *
 * Only for untyped connections.
 */
class DropQuery extends Query {
	// Table name.
	SStr table;

	// Create.
	init(SrcPos pos, SStr table) {
		init(pos) {
			table = table;
		}
	}

	// Resolve.
	void resolve(ResolveContext context) : override {
		if (context.typed)
			throw SyntaxError(pos, "Can not use CREATE TABLE queries with typed connections.");
	}

	// Build the query.
	QueryBuilder build() : override {
		QueryBuilder r;

		r.query.put("DROP TABLE ");
		r.query.name(table.v);
		r.query.put(";");

		r;
	}

}


/**
 * Custom syntax for transactions.
 */
class TransactionBlock extends Block {
	// Child block for actual code.
	ExprBlock child;

	// The variable created for the transaction object.
	private Var transactionVar;

	// Current block and label to exit the block.
	private core:asm:Label exitCommitLabel;
	private core:asm:Label exitRollbackLabel;
	private core:asm:Block exitBlock;
	private CodeResult exitResult;

	// Keep track of the current return type for this block.
	private ExprResult abortTypes;

	// Create, use an explicit expression for the connection.
	init(SrcPos pos, Block parent, Expr connection) {
		super(pos, parent);

		var result = connection.result.type;
		var toUse = if (Value(named{DBConnection}).mayReferTo(result)) {
			connection;
		} else if (result.type as DatabaseType) {
			MemberVarAccess(connection.pos, connection, named{TypedConnection:connection<TypedConnection>});
		} else {
			throw SyntaxError(connection.pos, "Expected a database connection (typed or untyped).");
		};

		init {
			child(pos, this);
			transactionVar(this, named{Transaction}, SStr(" transaction"), Actuals(toUse));
			abortTypes = noReturn;
		}
	}

	// Create, use a surrounding block.
	init(SrcPos pos, Block parent) {
		self(pos, parent, QueryBlock:find(parent).connection);
	}

	// Result type.
	ExprResult result() : override {
		var cResult = child.result;

		if (abortTypes.nothing) {
			cResult;
		} else if (cResult.nothing) {
			abortTypes;
		} else {
			var c = common(abortTypes.type, cResult.type);
			// Note: We explicitly allow using ROLLBACK or COMMIT to force the result to void.
			if (c == Value() & abortTypes.value) {
				StrBuf msg;
				msg << "Type mismatch in returns values from transaction. The early abort values result in "
					<< abortTypes.type << " while the final statement in the block results in "
					<< cResult.type << ", which is not compatible.";
				throw SyntaxError(pos, msg.toS);
			}
			c;
		}
	}

	// Called to "register" a new result from ROLLBACK or COMMIT stmts.
	void newAbortResult(SrcPos pos, ExprResult result) {
		if (result.nothing)
			return;

		if (abortTypes.nothing) {
			// First time.
			abortTypes = result;
		} else {
			var c = common(abortTypes.type, result.type);
			if (c == Value()) {
				StrBuf msg;
				msg << "Type mismatch in early abort values for transaction. ";
				msg << "So far, the common result type is " << abortTypes.type;
				msg << " but this can not be combined with " << result.type << ".";
				throw SyntaxError(pos, msg.toS);
			}
			abortTypes = c;
		}
	}

	// Generate code.
	void blockCode(CodeGen to, CodeResult result) : override {
		exitBlock = to.block;
		exitCommitLabel = to.l.label;
		exitRollbackLabel = to.l.label;
		var end = to.l.label;

		exitResult = result.split(to);

		transactionVar.code(to, CodeResult());
		child.code(to, exitResult);

		// Make sure it is activated properly.
		result.created(to);

		// Generate code to commit the transaction.
		to.l << exitCommitLabel;
		to.l << core:asm:lea(core:asm:ptrA, transactionVar.var.var.v);
		to.l << core:asm:fnParam(core:asm:ptrDesc, core:asm:ptrA);
		to.l << core:asm:fnCall(named{Transaction:commit<Transaction>}.ref, true);
		to.l << core:asm:jmp(end);

		// Generate code to roll back the transaction (so that we can jump there easily).
		to.l << exitRollbackLabel;
		to.l << core:asm:lea(core:asm:ptrA, transactionVar.var.var.v);
		to.l << core:asm:fnParam(core:asm:ptrDesc, core:asm:ptrA);
		to.l << core:asm:fnCall(named{Transaction:rollback<Transaction>}.ref, true);

		to.l << end;
	}

	// Either generate a commit or a rollback, optionally returning a value.
	void prematureAbort(CodeGen to, Bool commit, Expr? result) {
		if (result) {
			result.code(to, exitResult);
		}

		var lbl = if (commit) {
			exitCommitLabel;
		} else {
			exitRollbackLabel;
		};

		to.l << core:asm:jmpBlock(lbl, exitBlock);
	}

	// Helper to find a transaction.
	TransactionBlock find(Block from) : static {
		NameLookup? at = from.lookup;
		while (l = at as BlockLookup) {
			if (q = l.block as TransactionBlock)
				return q;

			at = l.parent;
		}

		throw InternalError("Can not find a transaction block.");
	}

	// To string.
	void toS(StrBuf to) : override {
		to << "TRANSACTION ";
		to << child;
	}
}


/**
 * Abort transactions prematurely.
 */
class TransactionAbort extends Expr {
	// Transaction to abort.
	private TransactionBlock transaction;

	// Value, if any.
	private Expr? resultExpr;

	// Commit?
	private Bool commit;

	// Create.
	init(SrcPos pos, Block parent, Bool commit, Expr? result) {
		init(pos) {
			transaction = TransactionBlock:find(parent);
			resultExpr = result;
			commit = commit;
		}

		if (result)
			transaction.newAbortResult(pos, result.result);
		else
			transaction.newAbortResult(pos, Value());
	}

	ExprResult result() : override {
		noReturn;
	}

	void code(CodeGen to, CodeResult r) : override {
		transaction.prematureAbort(to, commit, resultExpr);
	}

	void toS(StrBuf to) : override {
		if (commit)
			to << "COMMIT";
		else
			to << "ROLLBACK";
	}
}