swaf/src/db/ModelQuery.ts

import {Connection} from "mysql";

import {Pagination} from "../common/Pagination.js";
import Model from "./Model.js";
import ModelFactory from "./ModelFactory.js";
import ModelRelation, {RelationDatabaseProperties} from "./ModelRelation.js";
import {isQueryVariable, query, QueryResult, QueryVariable} from "./MysqlConnectionManager.js";


export default class ModelQuery<M extends Model> implements WhereFieldConsumer<M> {
    public static select<M extends Model>(factory: ModelFactory<M>, ...fields: QueryFields): ModelQuery<M> {
        fields = fields.map(v => v === '' ? new SelectFieldValue('none', 1, true) : v);
        return new ModelQuery(QueryType.SELECT, factory, fields.length > 0 ? fields : ['*']);
    }

    public static insert<M extends Model>(factory: ModelFactory<M>, data: Pick<M, keyof M>): ModelQuery<M> {
        const fields = [];
        for (const key of Object.keys(data)) {
            fields.push(new FieldValue(key, data[key], false));
        }
        return new ModelQuery(QueryType.INSERT, factory, fields);
    }

    public static update<M extends Model>(factory: ModelFactory<M>, data: Pick<M, keyof M>): ModelQuery<M> {
        const fields = [];
        for (const key of Object.keys(data)) {
            fields.push(new FieldValue(inputToFieldOrValue(key, factory.table), data[key], false));
        }
        return new ModelQuery(QueryType.UPDATE, factory, fields);
    }

    public static delete<M extends Model>(factory: ModelFactory<M>): ModelQuery<M> {
        return new ModelQuery(QueryType.DELETE, factory);
    }

    private readonly type: QueryType;
    private readonly factory: ModelFactory<M>;
    private readonly table: string;
    private readonly fields: QueryFields;
    private _leftJoin?: string;
    private _leftJoinAlias?: string;
    private _leftJoinOn: WhereFieldValue[] = [];
    private _where: (WhereFieldValue | WhereFieldValueGroup)[] = [];
    private _limit?: number;
    private _offset?: number;
    private _sortBy?: string;
    private _sortDirection?: 'ASC' | 'DESC';
    private readonly relations: string[] = [];
    private readonly subRelations: { [relation: string]: string[] | undefined } = {};
    private _pivot?: string[];
    private _union?: ModelQueryUnion;
    private _recursiveRelation?: RelationDatabaseProperties;
    private _reverseRecursiveRelation?: boolean;

    private constructor(type: QueryType, factory: ModelFactory<M>, fields?: QueryFields) {
        this.type = type;
        this.factory = factory;
        this.table = factory.table;
        this.fields = fields || [];
    }

    public leftJoin(table: string, alias?: string): this {
        this._leftJoin = table;
        this._leftJoinAlias = alias;
        return this;
    }

    public on(
        field1: string,
        field2: string,
        test: WhereTest = WhereTest.EQ,
        operator: WhereOperator = WhereOperator.AND,
    ): this {
        this._leftJoinOn.push(new WhereFieldValue(
            inputToFieldOrValue(field1), inputToFieldOrValue(field2), true, test, operator,
        ));
        return this;
    }

    public where(
        field: string,
        value: ModelFieldData,
        test: WhereTest = WhereTest.EQ,
        operator: WhereOperator = WhereOperator.AND,
    ): this {
        this._where.push(new WhereFieldValue(field, value, false, test, operator));
        return this;
    }

    public groupWhere(
        setter: (query: WhereFieldConsumer<M>) => void,
        operator: WhereOperator = WhereOperator.AND,
    ): this {
        this._where.push(new WhereFieldValueGroup(this.collectWheres(setter), operator));
        return this;
    }

    private collectWheres(setter: (query: WhereFieldConsumer<M>) => void): (WhereFieldValue | WhereFieldValueGroup)[] {
        // eslint-disable-next-line @typescript-eslint/no-this-alias
        const query = this;
        const wheres: (WhereFieldValue | WhereFieldValueGroup)[] = [];
        setter({
            where(
                field: string,
                value: ModelFieldData,
                test: WhereTest = WhereTest.EQ,
                operator: WhereOperator = WhereOperator.AND,
            ) {
                wheres.push(new WhereFieldValue(field, value, false, test, operator));
                return this;
            },
            groupWhere(
                setter: (query: WhereFieldConsumer<M>) => void,
                operator: WhereOperator = WhereOperator.AND,
            ) {
                wheres.push(new WhereFieldValueGroup(query.collectWheres(setter), operator));
                return this;
            },
        });
        return wheres;
    }

    public limit(limit: number, offset: number = 0): this {
        this._limit = limit;
        this._offset = offset;
        return this;
    }

    public sortBy(field: string, direction: SortDirection = 'ASC', raw: boolean = false): this {
        this._sortBy = raw ? field : inputToFieldOrValue(field);
        this._sortDirection = direction;
        return this;
    }

    /**
     * @param relations The relations field names to eagerload. To load nested relations, separate fields with '.'
     * (i.e.: "author.roles.permissions" loads authors, their roles, and the permissions of these roles)
     */
    public with(...relations: string[]): this {
        relations.forEach(relation => {
            const parts = relation.split('.');
            if (this.relations.indexOf(parts[0]) < 0) this.relations.push(parts[0]);
            if (parts.length > 1) {
                if (!this.subRelations[parts[0]]) this.subRelations[parts[0]] = [];
                this.subRelations[parts[0]]?.push(parts.slice(1).join('.'));
            }
        });
        return this;
    }

    public pivot(...fields: string[]): this {
        this._pivot = fields;
        return this;
    }

    public union(
        query: ModelQuery<Model>,
        sortBy: string, direction: SortDirection = 'ASC',
        raw: boolean = false,
        limit?: number,
        offset?: number,
    ): this {
        if (this.type !== QueryType.SELECT) throw new Error('Union queries are only implemented with SELECT.');

        this._union = {
            query: query,
            sortBy: raw ? sortBy : inputToFieldOrValue(sortBy),
            direction: direction,
            limit: limit,
            offset: offset,
        };
        return this;
    }

    public recursive(relation: RelationDatabaseProperties, reverse: boolean): this {
        if (this.type !== QueryType.SELECT) throw new Error('Recursive queries are only implemented with SELECT.');
        this._recursiveRelation = relation;
        this._reverseRecursiveRelation = reverse;
        return this;
    }

    public toString(final: boolean = false): string {
        let query = '';

        if (this._pivot) this.fields.push(...this._pivot);

        // Prevent wildcard and fields from conflicting
        const fields = this.fields.map(f => {
            const field = f.toString();
            if (field.startsWith('(')) return f; // Skip sub-queries
            return inputToFieldOrValue(field, this.table);
        }).join(',');

        let join = '';
        if (this._leftJoin) {
            join = ` LEFT JOIN \`${this._leftJoin}\``
                + (this._leftJoinAlias ? ` AS \`${this._leftJoinAlias}\`` : '')
                + ` ON ${this._leftJoinOn[0]}`;
            for (let i = 1; i < this._leftJoinOn.length; i++) {
                join += this._leftJoinOn[i].toString(false);
            }
        }

        let where = '';
        if (this._where.length > 0) {
            where = ` WHERE ${this._where[0]}`;
            for (let i = 1; i < this._where.length; i++) {
                where += this._where[i].toString(false);
            }
        }

        let limit = '';
        if (typeof this._limit === 'number') {
            limit = ` LIMIT ${this._limit}`;
            if (typeof this._offset === 'number' && this._offset !== 0) {
                limit += ` OFFSET ${this._offset}`;
            }
        }

        let orderBy = '';
        if (typeof this._sortBy === 'string') {
            orderBy = ` ORDER BY ${this._sortBy} ${this._sortDirection}`;
        }

        const table = `\`${this.table}\``;

        switch (this.type) {
            case QueryType.SELECT:
                if (this._recursiveRelation) {
                    const cteFields = fields.replace(RegExp(`${table}`, 'g'), 'o');
                    const idKey = this._reverseRecursiveRelation ?
                        this._recursiveRelation.foreignKey :
                        this._recursiveRelation.localKey;
                    const sortOrder = this._reverseRecursiveRelation ? 'DESC' : 'ASC';

                    query = `WITH RECURSIVE cte AS (`
                        + `SELECT ${fields},1 AS __depth, CONCAT(\`${idKey}\`) AS __path FROM ${table}${where}`
                        + ` UNION `
                        + `SELECT ${cteFields},c.__depth + 1,CONCAT(c.__path,'/',o.\`${idKey}\`) AS __path FROM ${table} AS o, cte AS c WHERE o.\`${this._recursiveRelation.foreignKey}\`=c.\`${this._recursiveRelation.localKey}\``
                        + `) SELECT * FROM cte${join}${orderBy || ` ORDER BY __path ${sortOrder}`}${limit}`;
                } else {
                    query = `SELECT ${fields} FROM ${table}${join}${where}${orderBy}${limit}`;
                }

                if (this._union) {
                    const unionOrderBy = this._union.sortBy ? ` ORDER BY ${this._union.sortBy} ${this._union.direction}` : '';
                    const unionLimit = typeof this._union.limit === 'number' ? ` LIMIT ${this._union.limit}` : '';
                    const unionOffset = typeof this._union.offset === 'number' ? ` OFFSET ${this._union.offset}` : '';
                    query = `(${query}) UNION ${this._union.query.toString(false)}${unionOrderBy}${unionLimit}${unionOffset}`;
                }
                break;
            case QueryType.INSERT: {
                const insertFields = this.fields.filter(f => f instanceof FieldValue)
                    .map(f => f as FieldValue);
                const insertFieldNames = insertFields.map(f => f.fieldName).join(',');
                const insertFieldValues = insertFields.map(f => f.fieldValue).join(',');
                query = `INSERT INTO ${table} (${insertFieldNames}) VALUES(${insertFieldValues})`;
                break;
            }
            case QueryType.UPDATE:
                query = `UPDATE ${table} SET ${fields}${where}${orderBy}${limit}`;
                break;
            case QueryType.DELETE:
                query = `DELETE FROM ${table}${where}${orderBy}${limit}`;
                break;

        }

        return final ? query : `(${query})`;
    }

    public build(): string {
        return this.toString(true);
    }

    public get variables(): QueryVariable[] {
        const variables: QueryVariable[] = [];
        this.fields.filter(v => v instanceof FieldValue)
            .flatMap(v => (v as FieldValue).variables)
            .forEach(v => variables.push(v));
        this._where.flatMap(v => this.getVariables(v))
            .forEach(v => variables.push(v));
        this._union?.query.variables.forEach(v => variables.push(v));
        return variables;
    }

    private getVariables(where: WhereFieldValue | WhereFieldValueGroup): QueryVariable[] {
        return where instanceof WhereFieldValueGroup ?
            where.fields.flatMap(v => this.getVariables(v)) :
            where.variables;
    }

    public async execute(connection?: Connection): Promise<QueryResult> {
        return await query(this.build(), this.variables, connection);
    }

    public async get(connection?: Connection): Promise<ModelQueryResult<M>> {
        const queryResult = await this.execute(connection);
        const models: ModelQueryResult<M> = [];
        models.originalData = [];

        if (this._pivot) models.pivot = [];

        // Eager loading init
        const relationMap: { [p: string]: ModelRelation<Model, Model, Model | Model[] | null>[] } = {};
        for (const relation of this.relations) {
            relationMap[relation] = [];
        }

        for (const result of queryResult.results) {
            const modelData: Record<string, unknown> = {};
            for (const field of Object.keys(result)) {
                modelData[field.split('.')[1] || field] = result[field];
            }

            const model = this.factory.create(modelData as Pick<M, keyof M>, false);
            models.push(model);
            models.originalData.push(modelData);

            if (this._pivot && models.pivot) {
                const pivotData: Record<string, unknown> = {};
                for (const field of this._pivot) {
                    pivotData[field] = result[field.split('.')[1]];
                }
                models.pivot.push(pivotData);
            }

            // Eager loading init map
            for (const relation of this.relations) {
                if (model[relation] === undefined) throw new Error(`Relation ${relation} doesn't exist on ${model.constructor.name}.`);
                if (!(model[relation] instanceof ModelRelation)) throw new Error(`Field ${relation} is not a relation on ${model.constructor.name}.`);
                relationMap[relation].push(model[relation] as ModelRelation<Model, Model, Model | Model[] | null>);
            }
        }

        // Eager loading execute
        for (const relationName of this.relations) {
            const relations = relationMap[relationName];
            if (relations.length > 0) {
                const allModels = await relations[0].eagerLoad(relations, this.subRelations[relationName]);
                await Promise.all(relations.map(r => r.populate(allModels)));
            }
        }

        return models;
    }

    public async paginate(page: number, perPage: number, connection?: Connection): Promise<ModelQueryResult<M>> {
        this.limit(perPage, (page - 1) * perPage);
        const result = await this.get(connection);
        result.pagination = new Pagination(page, perPage, await this.count(true, connection));
        return result;
    }

    public async first(): Promise<M | null> {
        const models = await this.limit(1).get();
        return models.length > 0 ? models[0] : null;
    }

    public async count(removeLimit: boolean = false, connection?: Connection): Promise<number> {
        if (removeLimit) {
            this._limit = undefined;
            this._offset = undefined;
        }
        this._sortBy = undefined;
        this._sortDirection = undefined;

        this.fields.splice(0, this.fields.length);
        this.fields.push(new SelectFieldValue('_count', 'COUNT(*)', true));

        const queryResult = await this.execute(connection);
        return Number(queryResult.results[0]['_count']);
    }
}

function inputToFieldOrValue(input: string, addTable?: string): string {
    if (input.startsWith('`') || input.startsWith('"') || input.startsWith("'")) {
        return input;
    }

    let parts = input.split('.');
    if (addTable && parts.length === 1) parts = [addTable, input]; // Add table disambiguation

    return parts.map(v => v === '*' ? v : `\`${v}\``).join('.');
}

export interface ModelQueryResult<M extends Model> extends Array<M> {
    originalData?: Record<string, unknown>[];
    pagination?: Pagination;
    pivot?: Record<string, unknown>[];
}

export enum QueryType {
    SELECT,
    INSERT,
    UPDATE,
    DELETE,
}

export enum WhereOperator {
    AND = 'AND',
    OR = 'OR',
}

export enum WhereTest {
    EQ = '=',
    NE = '!=',
    GT = '>',
    GE = '>=',
    LT = '<',
    LE = '<=',
    IN = ' IN ',
}

class FieldValue {
    protected readonly field: string;
    protected value: ModelFieldData;
    protected raw: boolean;

    public constructor(field: string, value: ModelFieldData, raw: boolean) {
        this.field = field;
        this.value = value;
        this.raw = raw;
    }

    public toString(first: boolean = true): string {
        return `${first ? '' : ','}${this.fieldName}${this.test}${this.fieldValue}`;
    }

    protected get test(): string {
        return '=';
    }

    public get variables(): QueryVariable[] {
        if (this.value instanceof ModelQuery) return this.value.variables;
        if (this.raw || this.value === null || this.value === undefined ||
            typeof this.value === 'boolean') return [];
        if (Array.isArray(this.value)) return this.value.map(value => {
            if (!isQueryVariable(value)) value = value.toString();
            return value;
        }) as QueryVariable[];

        let value = this.value;
        if (!isQueryVariable(value)) value = value.toString();
        return [value as QueryVariable];
    }

    public get fieldName(): string {
        return inputToFieldOrValue(this.field);
    }

    public get fieldValue(): ModelFieldData {
        let value: string;
        if (this.value instanceof ModelQuery) {
            value = this.value.toString(false);
        } else if (this.value === null || this.value === undefined) {
            value = 'null';
        } else if (typeof this.value === 'boolean') {
            value = String(this.value);
        } else if (this.raw) {
            value = this.value.toString();
        } else {
            value = Array.isArray(this.value) ?
                `(${'?'.repeat(this.value.length).split('').join(',')})` :
                '?';
        }
        return value;
    }
}

export class SelectFieldValue extends FieldValue {
    public toString(): string {
        let value: string;
        if (this.value instanceof ModelQuery) {
            value = this.value.toString(true);
        } else if (this.value === null || this.value === undefined) {
            value = 'null';
        } else if (typeof this.value === 'boolean') {
            value = String(this.value);
        } else {
            value = this.raw ?
                this.value.toString() :
                '?';
        }
        return `(${value}) AS \`${this.field}\``;
    }
}

class WhereFieldValue extends FieldValue {
    private readonly _test: WhereTest;
    private readonly operator: WhereOperator;

    public constructor(field: string, value: ModelFieldData, raw: boolean, test: WhereTest, operator: WhereOperator) {
        super(field, value, raw);
        this._test = test;
        this.operator = operator;
    }

    public toString(first: boolean = true): string {
        return (!first ? ` ${this.operator} ` : '') + super.toString(true);
    }

    protected get test(): string {
        if (this.value === null || this.value === undefined) {
            if (this._test === WhereTest.EQ) {
                return ' IS ';
            } else if (this._test === WhereTest.NE) {
                return ' IS NOT ';
            }
        }
        return this._test;
    }
}

class WhereFieldValueGroup {
    public readonly fields: (WhereFieldValue | WhereFieldValueGroup)[];
    public readonly operator: WhereOperator;

    public constructor(fields: (WhereFieldValue | WhereFieldValueGroup)[], operator: WhereOperator) {
        this.fields = fields;
        this.operator = operator;
    }

    public toString(first: boolean = true): string {
        let str = `${first ? '' : ` ${this.operator} `}(`;
        let firstField = true;
        for (const field of this.fields) {
            str += field.toString(firstField);
            firstField = false;
        }
        str += ')';
        return str;
    }
}

export interface WhereFieldConsumer<M extends Model> {
    where(field: string, value: ModelFieldData, test?: WhereTest, operator?: WhereOperator): this;

    groupWhere(setter: (query: WhereFieldConsumer<M>) => void, operator?: WhereOperator): this;
}

export type QueryFields = (string | SelectFieldValue | FieldValue)[];

export type SortDirection = 'ASC' | 'DESC';

type ModelQueryUnion = {
    query: ModelQuery<Model>,
    sortBy: string,
    direction: SortDirection,
    limit?: number,
    offset?: number,
};

export type ModelFieldData =
    | QueryVariable
    | ModelQuery<Model>
    | { toString(): string }
    | (QueryVariable | { toString(): string })[];