import { afterEach } from 'node:test';

import { db } from '../db/client.js';

type QueryCall = {
  text: string;
  params: unknown[] | undefined;
};

type MockQueryResult = {
  rowCount?: number;
  rows?: unknown[];
};

type MockHandler = MockQueryResult | ((call: QueryCall) => MockQueryResult | Promise<MockQueryResult>);

const originalQuery = db.query.bind(db);

export const mockDb = (...handlers: MockHandler[]) => {
  const calls: QueryCall[] = [];
  const queue = [...handlers];

  const mockedQuery = async (text: string, params?: unknown[]) => {
    const call = { text, params };
    calls.push(call);

    const next = queue.shift();

    if (!next) {
      throw new Error(`Unexpected query: ${text}`);
    }

    const result = typeof next === 'function' ? await next(call) : next;

    return {
      rowCount: result.rowCount ?? result.rows?.length ?? 0,
      rows: result.rows ?? [],
    };
  };

  (db as typeof db & {
    query: typeof db.query;
  }).query = mockedQuery as typeof db.query;

  return { calls };
};

afterEach(() => {
  (db as typeof db & {
    query: typeof originalQuery;
  }).query = originalQuery;
});
