import { Pool, type PoolClient, type QueryResultRow } from 'pg';

export type DatabaseTransaction = Pick<PoolClient, 'query'>;

export class Database {
    public constructor(private readonly pool: Pool) {}

    public async transaction<T>(
        work: (transaction: DatabaseTransaction) => Promise<T>,
    ): Promise<T> {
        const client = await this.pool.connect();
        try {
            await client.query('BEGIN');
            // atlas_app_login is NOINHERIT by design; application DML always uses the least-privileged group role.
            await client.query('SET LOCAL ROLE atlas_app');
            const result = await work(client);
            await client.query('COMMIT');
            return result;
        } catch (error) {
            await client.query('ROLLBACK');
            throw error;
        } finally {
            client.release();
        }
    }

    /** Every tenant-scoped query must run through this method. */
    public async workspaceTransaction<T>(
        workspaceId: string,
        work: (transaction: DatabaseTransaction) => Promise<T>,
    ): Promise<T> {
        return this.transaction(async (transaction) => {
            // SET LOCAL is scoped to this transaction and cannot leak a tenant id through the pool.
            await transaction.query(
                "SELECT set_config('app.workspace_id', $1, true)",
                [workspaceId],
            );
            return work(transaction);
        });
    }

    public async query<Row extends QueryResultRow = QueryResultRow>(
        text: string,
        values?: unknown[],
    ) {
        // Even non-tenant reads use the restricted role; callers cannot accidentally use the login role's privileges.
        return this.transaction((transaction) =>
            transaction.query<Row>(text, values),
        );
    }

    public async close() {
        await this.pool.end();
    }
}

export function createApplicationDatabase(connectionString: string): Database {
    return new Database(new Pool({ connectionString, max: 10 }));
}