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 是 NOINHERIT 登录角色；应用 DML 必须降为最小权限的 atlas_app。
            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();
        }
    }

    /** 每个带工作空间边界的业务读写都必须通过此方法执行。 */
    public async workspaceTransaction<T>(
        workspaceId: string,
        work: (transaction: DatabaseTransaction) => Promise<T>,
    ): Promise<T> {
        return this.transaction(async (transaction) => {
            // SET LOCAL 仅作用于当前事务，连接归还池后不会泄漏租户 ID。
            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[],
    ) {
        // 即使不带租户上下文的查询也使用受限角色，调用方不会误用登录角色权限。
        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 }));
}