vore-scenario-generator/src/server/db/querytypes.ts

export type QueryDefinition<T extends string> = {
	readonly parameters: {
		readonly [key in T]: { readonly index: number, readonly validator: (value: undefined) => string | number | null }
	}
	readonly query: string
	readonly output: (result: D1Result<object>) => unknown
}
export type QueryParameters<DefinitionT extends QueryDefinition<any>> = DefinitionT extends QueryDefinition<infer ParametersT> ? {
	readonly [key in ParametersT]: Exclude<Parameters<DefinitionT['parameters'][key]['validator']>[0], undefined>
} : never
export type BoundQuery<ResultT> = {
	readonly statement: D1PreparedStatement,
	readonly transformer: (result: D1Result<object>) => ResultT
}
export type QueryOutput<DefinitionT extends QueryDefinition<any>> = ReturnType<DefinitionT["output"]>
export type PreparedQuery<DefinitionT extends QueryDefinition<any>> = (values: QueryParameters<DefinitionT>) => BoundQuery<QueryOutput<DefinitionT>>
export type QueryDefinitions = { readonly [key: string]: QueryDefinition<any> }
export type PreparedQueries<T extends QueryDefinitions> = { readonly [key in keyof T]: PreparedQuery<T[key]> }

const QUERY_PARAM_HEURISTIC = /\?(\d+)/g;

function parameterIndexes(parameters: QueryDefinition<any>["parameters"]): Set<number> {
	let result = new Set<number>();
	for (const key of Object.keys(parameters)) {
		const value = parameters[key].index;
		if (result.has(value)) {
			throw Error(`found duplicate index ${value}`)
		}
		result.add(value)
	}
	return result;
}

function queryBindingIndexes(query: string): Set<number> {
	const result = new Set<number>()
	for (const binding of query.matchAll(QUERY_PARAM_HEURISTIC)) {
		result.add(parseInt(binding[1]))
	}
	return result
}

export function validatedDefinition<T extends QueryDefinition<any>>(definition: T): T {
	const queryBindings = queryBindingIndexes(definition.query);
	const parameters = parameterIndexes(definition.parameters)
	const missing = Array.from(queryBindings).filter(v => !parameters.has(v))
	const extra = Array.from(parameters).filter(v => !queryBindings.has(v))
	if (missing.length + extra.length > 0) {
		if (missing.length > 0) {
			if (extra.length > 0) {
				throw Error(`missing definitions for ${missing.map(v => `?${v}`).join(', ')} and don't need definitions for ${extra.map(v => `?${v}`).join(', ')}`)
			} else {
				throw Error(`missing definitions for ${missing.map(v => `?${v}`).join(', ')} `)
			}
		} else {
			throw Error(`don't need definitions for ${extra.map(v => `?${v}`).join(', ')}`)
		}
	}
	return definition;
}

export function validatedDefinitions<T extends QueryDefinitions>(definitions: T): T {
	for (const key of Object.keys(definitions) as (keyof T & string)[]) {
		try {
			validatedDefinition(definitions[key]);
		} catch (e) {
			throw Error(`when validating definition for ${key}: ${e}`)
		}
	}
	return definitions;
}

export function prepareQuery<T extends QueryDefinition<any>>(database: D1Database, definition: T): PreparedQuery<T> {
	const preparedStatement = database.prepare(definition.query);
	return function(values: QueryParameters<T>) {
		const bindings: unknown[] = new Array(Array.from(parameterIndexes(definition.parameters)).reduce((a, b) => Math.max(a, b), 0));
		for (const key of Object.keys(definition.parameters)) {
			bindings[definition.parameters[key].index - 1] = definition.parameters[key].validator(values[key]);
		}
		return {
			statement: preparedStatement.bind(...bindings),
			transformer: definition.output
		};
	} as PreparedQuery<T>;
}

export function prepareAllQueries<T extends QueryDefinitions>(database: D1Database, q: T): PreparedQueries<T> {
	const result: Partial<PreparedQueries<T>> = {};
	for (const key of Object.keys(q) as (keyof T & string)[]) {
		try {
			result[key] = prepareQuery(database, q[key])
		} catch (e) {
			throw Error(`when preparing ${key}: ${e}`)
		}
	}
	return result as PreparedQueries<T>;
}

export async function runQuery<T>(db: D1Database, query: BoundQuery<T>): Promise<T> {
	const startAt = performance.now()
	const [results] = await db.batch([query.statement]);
	const endAt = performance.now()
	console.info(`DB query time: ${endAt - startAt} / Runtime: ${results.meta.duration} / Rows read: ${results.meta.rows_read} / Rows written: ${results.meta.rows_written}`)
	return query.transformer(results as D1Result<object>);
}

export async function batchQueries<T extends [...unknown[]]>(db: D1Database, queries: { readonly [K in keyof T]: BoundQuery<T[K]> }): Promise<T> {
	const startAt = performance.now()
	const results = await db.batch(queries.map(q => q.statement));
	const endAt = performance.now()
	console.info(`DB transaction time: ${endAt - startAt} / Runtime: ${results.map(r => `${r.meta.duration ?? 0}`).join('+')} = ${results.reduce((t, r) => t + (r.meta.duration ?? 0), 0)} / Rows read: ${results.map(r => `${r.meta.rows_read ?? 0}`).join('+')} = ${results.reduce((t, r) => t + (r.meta.rows_read ?? 0), 0)} / Rows written: ${results.map(r => `${r.meta.rows_written ?? 0}`).join('+')} = ${results.reduce((t, r) => t + (r.meta.rows_written ?? 0), 0)}`)
	return results.map((result, index) => queries[index].transformer(result as D1Result<object>)) as T;
}

export class TypedDBWrapper {
	private readonly db: D1Database;

	constructor(db: D1Database) {
		this.db = db;
	}

	prepare<T extends QueryDefinition<any>>(query: T): PreparedQuery<T> {
		return prepareQuery(this.db, query);
	}

	prepareAll<T extends QueryDefinitions>(queries: T): PreparedQueries<T> {
		return prepareAllQueries(this.db, queries);
	}

	async run<T>(query: BoundQuery<T>): Promise<T> {
		return runQuery(this.db, query);
	}

	async batch<T extends [...unknown[]]>(...queries: { readonly [K in keyof T]: BoundQuery<T[K]>}): Promise<T> {
		return batchQueries<T>(this.db, queries);
	}
}

// TODO: Use the new run and batch functions to fix the Database class's methods