/*
 * Copyright (C) 2024-present Puter Technologies Inc.
 *
 * This file is part of Puter.
 *
 * Puter is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

import { readdirSync, readFileSync } from 'fs';
import { isAbsolute, resolve as resolvePath } from 'path';
import type { IConfig } from '../../types';
import { Span } from '../../util/span.js';
import { PuterClient } from '../types';
import { compareMigrationFilenames } from './migrationFilenames.js';

export interface WriteResult {
    insertId: number | bigint;
    affectedRows: number;
    anyRowsAffected: boolean;
    /**
     * Warnings the statement raised; nonzero means a value may have been stored
     * differently from how it was sent. Absent when the engine can't vouch that
     * a clean insert stored exactly what was sent, on a fresh id.
     */
    warningStatus?: number;
}

export interface BatchEntry {
    statement: string;
    values: unknown[];
}

type SqlJsonPath = readonly [string, ...string[]];

/** The slice of a pooled connection that migrations need. */
export interface MigrationConnection {
    query(sql: string): Promise<unknown>;
    release(): void;
}

/**
 * Base database client. Subclasses must override every method that throws here.
 *
 * Do not instantiate directly — use the factory exported from
 * `clients/database/index.ts` which picks the right implementation based on
 * `config.database.engine`.
 */
export class AbstractDatabaseClient extends PuterClient {
    /** Short name used by `case()` to pick engine-specific values. */
    readonly engineName: string = '';

    constructor(config: IConfig) {
        super(config);
    }

    // ------------------------------------------------------------------
    // Abstract interface — subclasses MUST override
    // ------------------------------------------------------------------

    /** Execute a read query. Returns an array of row objects. */
    async read(
        _query: string,
        _params: unknown[] = [],
    ): Promise<Record<string, unknown>[]> {
        throw new Error('DatabaseClient.read() not implemented');
    }

    /**
     * Read that prefers the primary database (useful when read-replicas may
     * have replication lag). In single-node setups this is identical to
     * `read()`.
     */
    async pread(
        _query: string,
        _params: unknown[] = [],
    ): Promise<Record<string, unknown>[]> {
        throw new Error('DatabaseClient.pread() not implemented');
    }

    /** Execute a write query (INSERT / UPDATE / DELETE). */
    async write(_query: string, _params: unknown[] = []): Promise<WriteResult> {
        throw new Error('DatabaseClient.write() not implemented');
    }

    /** Execute multiple write statements in a single transaction. */
    async batchWrite(_entries: BatchEntry[]): Promise<void> {
        throw new Error('DatabaseClient.batchWrite() not implemented');
    }

    // ------------------------------------------------------------------
    // Shared helpers (rely on the abstract methods above)
    // ------------------------------------------------------------------

    /**
     * Generate and execute an INSERT statement from a table name and a
     * key/value data object.
     */
    async insert(
        tableName: string,
        data: Record<string, unknown>,
    ): Promise<WriteResult> {
        const cols = Object.keys(data);
        const values = Object.values(data);
        const sql =
            `INSERT INTO ${this.quoteIdentifier(tableName)} ` +
            `(${cols.map((c) => this.quoteIdentifier(c)).join(', ')}) ` +
            `VALUES (${cols.map(() => '?').join(', ')})` +
            this.returningIdClause();
        return this.write(sql, values);
    }

    /** Whether `read()` goes to a separate replica rather than the primary. */
    protected hasReadReplica(): boolean {
        return false;
    }

    /**
     * The replica leg of `tryHardRead`. Override when `read()` falls back to
     * the primary by itself, or a replica failure reads the primary twice.
     */
    protected replicaRead(
        query: string,
        params: unknown[],
    ): Promise<Record<string, unknown>[]> {
        return this.read(query, params);
    }

    /**
     * Like `read()`, but an empty or failed replica read falls back to the
     * primary, so replication lag can't hide a row that was just written. By
     * default the primary read starts alongside the replica one; `race: false`
     * reads the primary only after the replica returns empty or fails. Without
     * a replica this is a single `read()`.
     */
    @Span('db.tryHardRead', (query: string) => ({ 'db.statement': query }))
    async tryHardRead(
        query: string,
        params: unknown[] = [],
        { race = true }: { race?: boolean } = {},
    ): Promise<Record<string, unknown>[]> {
        if (!this.hasReadReplica()) return this.read(query, params);

        const primary = race ? settle(this.pread(query, params)) : null;
        try {
            const rows = await this.replicaRead(query, params);
            if (rows.length > 0) return rows;
        } catch {
            // replica failed — fall through to primary
        }
        const result = await (primary ?? settle(this.pread(query, params)));
        if ('error' in result) throw result.error;
        return result.value;
    }

    /**
     * Apply the `.sql` files named `<engineName>*` in each configured migration
     * directory, in filename order. Files must be idempotent: nothing records
     * which ones already ran. A failure aborts startup. `transactional` wraps
     * each file in BEGIN/COMMIT, for engines whose DDL can roll back.
     */
    protected async applyMigrationPaths(
        connect: () => Promise<MigrationConnection>,
        splitStatements: (contents: string) => string[],
        { transactional }: { transactional: boolean },
    ): Promise<void> {
        const paths = this.config.database?.migrationPaths;
        if (!paths || paths.length === 0) return;

        const tag = `[${this.engineName}]`;
        const conn = await connect();
        try {
            for (const rawPath of paths) {
                const dir = isAbsolute(rawPath)
                    ? rawPath
                    : resolvePath(process.cwd(), rawPath);

                let files: string[];
                try {
                    files = readdirSync(dir)
                        .filter(
                            (f) =>
                                f.endsWith('.sql') &&
                                f.startsWith(this.engineName),
                        )
                        .sort(compareMigrationFilenames);
                } catch (e) {
                    throw new Error(
                        `${tag} migration path is unreadable: ${dir}`,
                        {
                            cause: e,
                        },
                    );
                }

                if (files.length === 0) {
                    console.log(`${tag} no migrations in ${dir}`);
                    continue;
                }

                console.log(
                    `${tag} running migrations from ${dir}: ${files.length} file(s)`,
                );

                for (const file of files) {
                    const contents = readFileSync(
                        resolvePath(dir, file),
                        'utf8',
                    );
                    const statements = splitStatements(contents);
                    if (transactional) await conn.query('BEGIN');
                    try {
                        for (let i = 0; i < statements.length; i++) {
                            try {
                                await conn.query(statements[i]);
                            } catch (e) {
                                throw new Error(
                                    `${tag} failed to apply ${file} at statement ${i}`,
                                    { cause: e },
                                );
                            }
                        }
                        if (transactional) await conn.query('COMMIT');
                    } catch (e) {
                        if (transactional) {
                            try {
                                await conn.query('ROLLBACK');
                            } catch {
                                // the original failure is the one worth surfacing
                            }
                        }
                        throw e;
                    }
                    console.log(
                        `${tag} applied ${file} (${statements.length} statements)`,
                    );
                }
            }
        } finally {
            conn.release();
        }
    }

    /**
     * Return the value from `choices` that matches the current engine.
     *
     * Usage:
     *
     *     db.case({
     *         sqlite: "datetime('now')",
     *         mysql: 'NOW()',
     *         otherwise: 'NOW()',
     *     });
     *
     * If the engine name isn't present in `choices`, falls back to
     * `choices.otherwise`.
     */
    case<T>(choices: Record<string, T> & { otherwise?: T }): T {
        if (Object.prototype.hasOwnProperty.call(choices, this.engineName)) {
            return choices[this.engineName];
        }
        return choices.otherwise as T;
    }

    quoteIdentifier(identifier: string): string {
        return identifier
            .split('.')
            .map((part) => {
                if (part === '*') return part;
                return `\`${part.replaceAll('`', '``')}\``;
            })
            .join('.');
    }

    booleanLiteral(value: boolean): string {
        return value ? '1' : '0';
    }

    booleanValue(value: boolean): boolean | 0 | 1 {
        return value ? 1 : 0;
    }

    insertIgnoreInto(tableName: string): string {
        const table = this.quoteIdentifier(tableName);
        return this.case({
            sqlite: `INSERT OR IGNORE INTO ${table}`,
            postgres: `INSERT INTO ${table}`,
            otherwise: `INSERT IGNORE INTO ${table}`,
        });
    }

    insertIgnoreSuffix(): string {
        return this.case({
            postgres: ' ON CONFLICT DO NOTHING',
            otherwise: '',
        });
    }

    upsertClause(
        conflictColumns: readonly string[],
        updateColumns: readonly string[],
    ): string {
        if (updateColumns.length === 0) {
            throw new Error('upsertClause requires at least one update column');
        }

        const updateList = updateColumns
            .map((column) => `${this.quoteIdentifier(column)} = ?`)
            .join(', ');

        return this.case({
            mysql: `ON DUPLICATE KEY UPDATE ${updateList}`,
            otherwise: `ON CONFLICT(${conflictColumns
                .map((column) => this.quoteIdentifier(column))
                .join(', ')}) DO UPDATE SET ${updateList}`,
        });
    }

    jsonTextExtract(jsonExpression: string, path: SqlJsonPath): string {
        const sqlitePath = `$${path.map((part) => `.${part}`).join('')}`;
        return this.case({
            sqlite: `json_extract(${jsonExpression}, ${this.sqlStringLiteral(sqlitePath)})`,
            mysql: `JSON_UNQUOTE(JSON_EXTRACT(${jsonExpression}, ${this.sqlStringLiteral(sqlitePath)}))`,
            postgres: `${jsonExpression} #>> ARRAY[${path
                .map((part) => this.sqlStringLiteral(part))
                .join(', ')}]`,
            otherwise: `JSON_UNQUOTE(JSON_EXTRACT(${jsonExpression}, ${this.sqlStringLiteral(sqlitePath)}))`,
        });
    }

    /**
     * Whole `periodSeconds` buckets since the epoch, from a unix-seconds
     * column. MySQL's `/` yields a decimal, which would split a bucket.
     */
    periodBucket(secondsExpression: string, periodSeconds: number): string {
        return this.case({
            mysql: `(${secondsExpression} DIV ${periodSeconds})`,
            otherwise: `(${secondsExpression} / ${periodSeconds})`,
        });
    }

    /**
     * Counts each (`userExpression`, bucket) pair once. `COUNT(DISTINCT a, b)`
     * is MySQL-only, so the pair folds into one number — injective while the
     * bucket stays under the multiplier, which hourly reaches in the 2980s.
     */
    countDistinctUserPeriods(
        userExpression: string,
        secondsExpression: string,
        periodSeconds: number,
    ): string {
        return `COUNT(DISTINCT ${userExpression} * 10000000 + ${this.periodBucket(secondsExpression, periodSeconds)})`;
    }

    nullCoalesce(...expressions: readonly string[]): string {
        if (expressions.length === 0) {
            throw new Error('nullCoalesce requires at least one expression');
        }
        return `COALESCE(${expressions.join(', ')})`;
    }

    returningIdClause(): string {
        return this.case({
            postgres: ' RETURNING id',
            otherwise: '',
        });
    }

    /**
     * A timestamp column value from one bound parameter in unix seconds, stored
     * the way `CURRENT_TIMESTAMP` would store that instant.
     */
    timestampFromUnixSql(): string {
        return this.case({
            sqlite: "datetime(?, 'unixepoch')",
            postgres: 'to_timestamp(?)',
            otherwise: 'FROM_UNIXTIME(?)',
        });
    }

    /**
     * A timestamp column as unix seconds, computed by the database: the mysql
     * driver reads a stored UTC timestamp as local time, so parsing it in JS
     * can be off by the server's UTC offset.
     */
    unixSecondsSql(column: string): string {
        return this.case({
            postgres: `EXTRACT(EPOCH FROM ${column})::bigint`,
            mysql: `UNIX_TIMESTAMP(${column})`,
            otherwise: `CAST(strftime('%s', ${column}) AS INTEGER)`,
        });
    }

    /**
     * What a read of a timestamp column holding `seconds` returns: sqlite gives
     * back the UTC text it stored, mysql2 and pg a `Date` read in node's time
     * zone (so this matches when node and the database session share one).
     */
    timestampValue(seconds: number): Date | string {
        const date = new Date(seconds * 1000);
        return this.case<Date | string>({
            sqlite: date.toISOString().slice(0, 19).replace('T', ' '),
            otherwise: date,
        });
    }

    protected sqlStringLiteral(value: string): string {
        return `'${value.replaceAll("'", "''")}'`;
    }
}

/** Await later without an unhandled rejection in the meantime. */
const settle = async <T>(
    promise: Promise<T>,
): Promise<{ value: T } | { error: unknown }> => {
    try {
        return { value: await promise };
    } catch (error) {
        return { error };
    }
};