import Fastify from 'fastify';
import { afterEach, describe, expect, it, vi } from 'vitest';
import { drainWithTimeout, registerGracefulShutdown } from '../src/shutdown.js';

function deferred<T = void>() {
  let resolve!: (value: T) => void;
  const promise = new Promise<T>((r) => {
    resolve = r;
  });
  return { promise, resolve };
}

describe('drainWithTimeout', () => {
  afterEach(() => {
    vi.useRealTimers();
  });

  it('resolves "drained" when the work settles first and clears its timer', async () => {
    vi.useFakeTimers();
    const work = deferred();
    const outcome = drainWithTimeout(work.promise, 30_000);
    work.resolve();
    await expect(outcome).resolves.toBe('drained');
    expect(vi.getTimerCount()).toBe(0);
  });

  it('resolves "timeout" when the deadline wins', async () => {
    vi.useFakeTimers();
    const never = new Promise<void>(() => undefined);
    const outcome = drainWithTimeout(never, 30_000);
    await vi.advanceTimersByTimeAsync(30_000);
    await expect(outcome).resolves.toBe('timeout');
  });
});

describe('registerGracefulShutdown', () => {
  it('on signal: stops the poller, drains the runner inside app.close(), and exits 0', async () => {
    const app = Fastify({ logger: false });
    const drain = deferred();
    const stop = deferred();
    const runner = { drain: vi.fn(() => drain.promise) };
    const poller = { stop: vi.fn(() => stop.promise) };
    const exited = deferred<number>();
    const shutdown = registerGracefulShutdown({ app, poller, runner, timeoutMs: 5000, exit: exited.resolve });
    await app.ready();

    shutdown.onSignal('SIGTERM');
    // Wait for the onClose hook to actually start the drain.
    await vi.waitFor(() => expect(runner.drain).toHaveBeenCalledTimes(1));
    expect(poller.stop).toHaveBeenCalledTimes(1);

    // Exit must wait for BOTH the runner drain and the in-flight poll tick.
    let exitCode: number | undefined;
    void exited.promise.then((code) => {
      exitCode = code;
    });
    drain.resolve();
    await new Promise((r) => setImmediate(r));
    expect(exitCode).toBeUndefined();

    stop.resolve();
    expect(await exited.promise).toBe(0);
    expect(shutdown.timedOut).toBe(false);
  });

  it('exits 1 when in-flight work does not settle within the timeout', async () => {
    const app = Fastify({ logger: false });
    const runner = { drain: () => new Promise<void>(() => undefined) }; // never settles
    const poller = { stop: () => Promise.resolve() };
    const exited = deferred<number>();
    const shutdown = registerGracefulShutdown({ app, poller, runner, timeoutMs: 50, exit: exited.resolve });
    await app.ready();

    shutdown.onSignal('SIGINT');
    expect(await exited.promise).toBe(1);
    expect(shutdown.timedOut).toBe(true);
  });

  it('exits 1 when close() itself rejects', async () => {
    const app = Fastify({ logger: false });
    const exited = deferred<number>();
    const shutdown = registerGracefulShutdown({
      app,
      poller: { stop: () => Promise.resolve() },
      runner: { drain: () => Promise.reject(new Error('drain exploded')) },
      timeoutMs: 5000,
      exit: exited.resolve,
    });
    await app.ready();

    shutdown.onSignal('SIGTERM');
    expect(await exited.promise).toBe(1);
  });
});