import { describe, it, expect, vi, beforeEach } from 'vitest'
import type { Client } from 'pg'
import { closePgClientSafely } from '@/lib/realtime/pg-client-cleanup'

function makeFakeClient(opts: {
  endResolves?: Promise<void>
  destroy?: ReturnType<typeof vi.fn>
}): Client {
  const handlers = new Map<string, Array<(...args: unknown[]) => void>>()
  const fake = {
    end: vi.fn().mockReturnValue(opts.endResolves ?? Promise.resolve()),
    on: vi.fn((event: string, fn: (...args: unknown[]) => void) => {
      const list = handlers.get(event) ?? []
      list.push(fn)
      handlers.set(event, list)
      return fake
    }),
    removeAllListeners: vi.fn((event: string) => {
      handlers.delete(event)
      return fake
    }),
    connection: {
      stream: { destroy: opts.destroy ?? vi.fn() },
    },
  }
  return fake as unknown as Client
}

describe('closePgClientSafely', () => {
  beforeEach(() => {
    vi.useRealTimers()
  })

  it('drops listeners and awaits client.end() when it resolves quickly', async () => {
    const destroy = vi.fn()
    const client = makeFakeClient({ destroy })

    await closePgClientSafely(client, 'test')

    expect(client.removeAllListeners).toHaveBeenCalledWith('notification')
    expect(client.removeAllListeners).toHaveBeenCalledWith('error')
    expect(client.end).toHaveBeenCalledOnce()
    expect(destroy).not.toHaveBeenCalled() // ended in time
  })

  it('falls back to socket-destroy when client.end() hangs past the timeout', async () => {
    const destroy = vi.fn()
    // .end() never resolves
    const client = makeFakeClient({ endResolves: new Promise(() => {}), destroy })

    vi.useFakeTimers()
    const promise = closePgClientSafely(client, 'test-hang')
    await vi.advanceTimersByTimeAsync(2_001)
    await promise

    expect(destroy).toHaveBeenCalledOnce()
    const arg = destroy.mock.calls[0][0]
    expect(arg).toBeInstanceOf(Error)
  })

  it('does not throw when client.end() rejects', async () => {
    const client = makeFakeClient({ endResolves: Promise.reject(new Error('boom')) })

    await expect(closePgClientSafely(client, 'test-reject')).resolves.toBeUndefined()
  })
})