import {AnyModel, IdentifierType, Model, ModelClass, ModelShape, SerializedModel} from "@sharkitek/core";
import {eq, getTableName, Table, TableConfig} from "drizzle-orm";
import {PgDatabase} from "drizzle-orm/pg-core";
import {PgQueryResultHKT} from "drizzle-orm/pg-core/session";
import type {DBQueryConfig, ExtractTablesWithRelations, TablesRelationalConfig} from "drizzle-orm/relations";
import type {KnownKeysOnly} from "drizzle-orm/utils";
import {DrizzleExtension} from "./Drizzle";
import {inArray} from "drizzle-orm/sql/expressions/conditions";

/**
 * Model query manager.
 */
export class ModelQuery<
	ModelType extends Model<Shape, IdentifierType<Shape, Identifier>>,
	TableType extends Table<TC>, TC extends TableConfig,
	TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown>,
	Shape extends ModelShape, Identifier extends keyof Shape = any, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>
>
{
	constructor(protected modelClass: ModelClass<ModelType, Shape, Identifier>, protected modelInstance: ModelType, protected database: PgDatabase<TQueryResult, TFullSchema, TSchema>, protected table: TableType)
	{}

	/**
	 * Parse the given model data to a model.
	 * @param rawModel Raw model.
	 * @protected
	 */
	protected parseModel(rawModel: SerializedModel<Shape>|null): ModelType|null
	{
		// The raw model is null or undefined, return NULL.
		if (!rawModel) return null;
		// Parse the given raw model to a model.
		return (new this.modelClass()).deserialize(rawModel) as ModelType;
	}

	/**
	 * Find models matching the given configuration.
	 * @param config Request configuration.
	 */
	async get<TConfig extends DBQueryConfig<"many", true, TSchema, TSchema[keyof TSchema]>>(config?: KnownKeysOnly<TConfig, DBQueryConfig<"many", true, TSchema, TSchema[keyof TSchema]>>): Promise<ModelType[]>
	{
		// Parse retrieved raw models.
		return (
			// Find many models from the given configuration.
			await ((this.database.query as any)?.[getTableName(this.table)]?.findMany(config) as Promise<SerializedModel<Shape>[]>)
		).map(rawModel => this.parseModel(rawModel)).filter((model) => !!model); // Only keep non-null models.
	}

	/**
	 * Find the first model matching the given configuration.
	 * @param config Request configuration.
	 */
	async first<TSelection extends Omit<DBQueryConfig<"many", true, TSchema, TSchema[keyof TSchema]>, "limit">>(config?: KnownKeysOnly<TSelection, Omit<DBQueryConfig<"many", true, TSchema, TSchema[keyof TSchema]>, "limit">>): Promise<ModelType|null>
	{
		// Parse retrieved raw model.
		return this.parseModel(
			// Find a model from the given configuration.
			await ((this.database.query as any)?.[getTableName(this.table)]?.findFirst(config) as Promise<SerializedModel<Shape>>)
		);
	}

	/**
	 * Find a model from its identifier.
	 * @param identifier Model identifier.
	 */
	async find(identifier: IdentifierType<Shape, Identifier>): Promise<ModelType|null>
	{
		// Find the first model which match the given identifier.
		return this.first({
			where: eq((this.table as any)?.[this.modelInstance.getIdentifierName()], identifier),
		});
	}

	/**
	 * Find many models from their identifier.
	 * @param identifiers Model identifiers.
	 */
	async findMany(...identifiers: (IdentifierType<Shape, Identifier>|IdentifierType<Shape, Identifier>[])[]): Promise<ModelType[]>
	{
		// Find many models which match the given identifiers.
		return this.get({
			where: inArray((this.table as any)?.[this.modelInstance.getIdentifierName()], identifiers.flat()),
		});
	}
}

/**
 * Model query builder.
 * @param modelClass Model class.
 */
export function query<
	ModelType extends Model<Shape, IdentifierType<Shape, Identifier>> & DrizzleExtension<AnyModel, TableType, TC, TQueryResult, TFullSchema, Scopes, any, any, TSchema>,
	TableType extends Table<TC>, TC extends TableConfig,
	TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown>,
	Scopes extends object,
	Shape extends ModelShape, Identifier extends keyof Shape = any, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>
>(modelClass: ModelClass<ModelType, Shape, Identifier>): ModelQuery<ModelType, TableType, TC, TQueryResult, TFullSchema, Shape, Identifier, TSchema> & Scopes
{
	// Get model drizzle instance.
	const model = new modelClass();
	const modelDrizzle = model.drizzle();
	// Create a model query with the model database and table.
	return Object.assign(
		new ModelQuery<ModelType, TableType, TC, TQueryResult, TFullSchema, Shape, Identifier, TSchema>(
			modelClass, model, modelDrizzle.getDatabase(), modelDrizzle.getTable(),
		),
		modelDrizzle.getScopes(),
	);
}