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

/**
 * Description of the columns in a table. Support nesting identifiers to an arbitrary depth, even if
 * we only ever use a single level.
 */
class TypedCol on Compiler {
	// Name.
	Str name;

	init(Str name) {
		init { name = name; }
	}

	Type varType() : abstract;

	Nat? firstNonNull() : abstract;

	void toS(StrBuf to) : override {
		to << name;
	}
}

/**
 * Single column.
 */
class ScalarTypedCol extends TypedCol {
	Nat id;
	Type type;

	init(Str name, Nat id, Type type) {
		init(name) { id = id; type = type; }
	}

	Type varType() : override {
		type;
	}

	Nat? firstNonNull() : override {
		if (isMaybe(Value(type)))
			return null;
		else
			return id;
	}

	void toS(StrBuf to) : override {
		super:toS(to);
		to << "->" << type.identifier << "@" << id;
	}
}

/**
 * Nested array of columns.
 *
 * Note: We need to preserve the order, so we can't use Map.
 */
class NestedTypedCol extends TypedCol {
	TypedCol[] columns;
	Bool maybe;

	init(Str name, Bool maybe) {
		init(name) { maybe = maybe; }
	}

	void push(TypedCol col) {
		columns.push(col);
	}

	Type varType() : override {
		Type r = TypedRow(columns);
		if (maybe) {
			unless (m = wrapMaybe(Value(r)).type)
				throw InternalError("Failed to find the maybe type.");
			r = m;
		}
		r;
	}

	Nat? firstNonNull() : override {
		for (c in columns)
			if (r = c.firstNonNull())
				return r;
		null;
	}

	void toS(StrBuf to) : override {
		super:toS(to);
		to << "->";
		if (maybe)
			to << "?";
		to << "(" << join(columns, ",") << ")";
	}
}



/**
 * Type used to represent a row from an SQL query.
 *
 * Might consist of two levels to represent 'table'.'column' properly, and to allow wrapping entire
 * sub-tables in Maybe types to avoid checking for null in case of left- or right- joins.
 */
class TypedRow extends Type {
	init(TypedCol[] columns) {
		init("<row>", TypeFlags:typeClass) {}

		parentLookup = named{};

		// Add member variables.
		MemberVar[] vars;
		for (col in columns) {
			MemberVar v(col.name, col.varType(), this);
			vars << v;
			add(v);
		}

		// Add a constructor that reads data from a Row instance.
		BSTreeCtor ctor([thisParam(this), ValParam(named{Row}, "row")], SrcPos());
		CtorBody body(ctor, Scope(this));

		unless (rowVar = body.variable(SimplePart("row")))
			throw InternalError("Failed to find the local variable 'row'.");
		LocalVarAccess row(SrcPos(), rowVar);

		{
			InitBlock init(SrcPos(), body, null);
			for (i, col in columns) {
				init.init(Initializer(SStr(col.name), getColumn(body, row, vars[i].type, col)));
			}
			body.add(init);
		}

		ctor.body = body;
		add(ctor);

		// To string.
		add(toString(vars));

		// Other default operators.
		add(TypeDeepCopy(this));
		add(TypeCopyCtor(this));
	}

	// Generate a toS function.
	private Function toString(MemberVar[] vars) {
		BSTreeFn fn(Value(), SStr("toS"), [thisParam(this), ValParam(named{StrBuf}, "to")], null);
		FnBody body(fn, Scope(this, BSLookup()));

		unless (toVar = body.variable(SimplePart("to")))
			throw InternalError("Failed to find local variable 'to'.");
		unless (thisVar = body.variable(SimplePart("this")))
			throw InternalError("Failed to find local variable 'this'.");
		LocalVarAccess to(SrcPos(), toVar);
		LocalVarAccess me(SrcPos(), thisVar);

		body.add(namedExpr(body, SrcPos(), "<<", to, Actuals(StrLiteral(SrcPos(), "{ "))));

		for (i, var in vars) {
			Str s = var.name + ": ";
			if (i > 0)
				s = ", " + s;
			body.add(namedExpr(body, SrcPos(), "<<", to, Actuals(StrLiteral(SrcPos(), s))));

			MemberVarAccess access(SrcPos(), me, var, true);
			Expr value = access;
			if (access.result.type.type !is named{Str})
				value = namedExpr(body, SrcPos(), "toS", access);
			body.add(namedExpr(body, SrcPos(), "<<", to, Actuals(value)));
		}

		body.add(namedExpr(body, SrcPos(), "<<", to, Actuals(StrLiteral(SrcPos(), " }"))));

		fn.body = body;
		fn;
	}

	// Generate an expression for extracting a particular column.
	private Expr getColumn(Block block, Expr row, Value type, TypedCol col) : static {
		if (col as ScalarTypedCol) {
			if (isMaybe(type)) {
				return getMaybeColumn(block, row, type, unwrapMaybe(type), col.id);
			} else {
				return getPlainColumn(block, row, type, col.id);
			}
		} else if (col as NestedTypedCol) {
			unless (nested = unwrapMaybe(type).type as TypedRow)
				throw InternalError("Nested type rows should be implemented as a TypedRow type!");

			if (col.maybe) {
				return getMaybeNested(block, row, type, nested, col);
			} else {
				return getPlainNested(block, row, nested);
			}
		}

		throw InternalError("Unknown subclass of TypedCol found!");
	}

	// Extract a column, given that we know it is not a maybe type.
	private Expr getPlainColumn(Block block, Expr row, Value type, Nat id) : static {
		var fn = getColumnFn(type.type);
		Actuals params;
		params.add(row);
		params.add(NumLiteral(SrcPos(), id.long));
		return FnCall(SrcPos(), Scope(), fn, params);
	}

	// Extract a column, but check for NULL first.
	private Expr getMaybeColumn(Block block, Expr row, Value maybeType, Value plainType, Nat id) : static {
		var fn = named{Row:isNull<Row, Nat>};
		Actuals params;
		params.add(row);
		params.add(NumLiteral(SrcPos(), id.long));
		If check(block, FnCall(SrcPos(), Scope(), fn, params));

		check.success(callCtor(check.successBlock, maybeType, Actuals()));
		check.fail(callCtor(block, maybeType, Actuals(getPlainColumn(block, row, plainType, id))));

		check;
	}

	// Extract a column, as a nested type.
	private Expr getPlainNested(Block block, Expr row, TypedRow type) : static {
		return callCtor(block, Value(type), Actuals(row));
	}

	// Extract a column, but check for NULL first.
	private Expr getMaybeNested(Block block, Expr row, Value wrapped, TypedRow plain, NestedTypedCol col) : static {
		// If all members are null, then we assume that the entire structure is never null since we
		// can't tell the difference between the entire table being null, and all the columns. We
		// *could* check all columns and assume that all nulls equal a null structure, but that is too
		// expensive for a heuristic that is likely seldom used (primary keys are usually not null, for
		// example).
		unless (check = col.firstNonNull)
			return callCtor(block, wrapped, Actuals(getPlainNested(block, row, plain)));

		var fn = named{Row:isNull<Row, Nat>};
		Actuals params;
		params.add(row);
		params.add(NumLiteral(SrcPos(), check.long));

		If check(block, FnCall(SrcPos(), Scope(), fn, params));
		check.success(callCtor(check.successBlock, wrapped, Actuals()));
		check.fail(callCtor(block, wrapped, Actuals(getPlainNested(block, row, plain))));
		check;
	}

	// Call a constructor.
	private Expr callCtor(Block block, Value type, Actuals actuals) : static {
		unless (t = type.type)
			throw InternalError("Could not find a proper type.");

		var params = actuals.values();
		params.insert(0, type);
		unless (fn = t.find("__init", params, Scope()) as Function)
			throw InternalError("Could not find a constructor for ${type}");

		return CtorCall(SrcPos(), Scope(), fn, actuals);
	}
}


/**
 * Type used to denote an iterator that produces instances of a TypedRow class.
 *
 * Follows the "simple" iterator type that only has a "next" function.
 */
class TypedIter extends Type {
	init(TypedCol[] types) {
		init("<iter>", TypeFlags:typeValue) {}

		parentLookup = named{};

		TypedRow toCreate = getTypedRow(types);
		MemberVar iter("base", Value(named{Statement:Result}), this);
		iter.visibility = typePrivate();
		add(iter);

		addCtor();
		addNext(toCreate, iter);
		addIter();
		add(TypeCopyCtor(this));
		add(TypeAssign(this));
		add(TypeDefaultDtor(this));
	}

	// Add constructor.
	private void addCtor() {
		BSTreeCtor ctor([thisParam(this), ValParam(named{Statement:Result}, "result")], SrcPos());
		CtorBody body(ctor, Scope(this, BSLookup()));

		unless (resultVar = body.variable(SimplePart("result")))
			throw InternalError("Failed to find local variable 'result'.");
		LocalVarAccess result(SrcPos(), resultVar);

		InitBlock init(SrcPos(), body, null);
		init.init(Initializer(SStr("base"), result));
		body.add(init);

		ctor.body = body;
		add(ctor);
	}

	// Add "next" member.
	private void addNext(TypedRow row, MemberVar iter) {
		BSTreeFn fn(wrapMaybe(Value(row)), SStr("next"), [thisParam(this)], null);
		FnBody body(fn, Scope(this, BSLookup()));

		unless (thisVar = body.variable(SimplePart("this")))
			throw InternalError("Failed to find local variable 'this'.");
		LocalVarAccess me(SrcPos(), thisVar);

		unless (typeCtor = row.find("__init", [Value(row), Value(named{Row})], Scope()) as Function)
			throw InternalError("Failed to find the constructor of the generated row-type.");

		// Call "next" and see if it returned null.
		WeakMaybeCast cast(namedExpr(body, SrcPos(), "next", MemberVarAccess(SrcPos(), me, iter, true)));
		cast.name(SStr("x"));
		If check(body, cast);
		body.add(check);

		if (created = cast.result) {
			CtorCall c(SrcPos(), Scope(), typeCtor, Actuals(LocalVarAccess(SrcPos(), created)));
			check.success(Return(SrcPos(), check.successBlock, c));
		}

		// Just return null otherwise.
		body.add(NullExpr(SrcPos()));

		fn.body = body;
		add(fn);
	}

	// Add "iter" member. We just return ourselves. Otherwise, we can't use it in foreach-loops. The
	// "clean" option would be to make a separate class that is "some returned rows" that can be
	// iterated. But this makes more sense overall.
	private void addIter() {
		BSTreeFn fn(Value(this), SStr("iter"), [thisParam(this)], null);
		FnBody body(fn, Scope(this));

		unless (thisVar = body.variable(SimplePart("this")))
			throw InternalError("Failed to find local variable 'this'.");
		LocalVarAccess me(SrcPos(), thisVar);
		body.add(me);

		fn.body = body;
		add(fn);
	}
}


// Global cache of data types.
private Map<Str, TypedRow> createdRowTypes on Compiler;
private Map<Str, TypedIter> createdIterTypes on Compiler;

// Compute a key for use in the map.
private Str typeKey(TypedCol[] types) on Compiler {
	StrBuf s;
	s << join(types, ",");
	return s.toS;
}

// Get a TypedRow instance.
TypedRow getTypedRow(TypedCol[] types) on Compiler {
	Str key = typeKey(types);
	if (rt = createdRowTypes.at(key))
		return rt;

	TypedRow created(types);
	createdRowTypes.put(key, created);
	return created;
}

// Get a TypedIter instance.
TypedIter getTypedIter(TypedCol[] types) on Compiler {
	Str key = typeKey(types);
	if (it = createdIterTypes.at(key))
		return it;

	TypedIter created(types);
	createdIterTypes.put(key, created);
	return created;
}