import type { FastMCP } from 'fastmcp';
import { ServerLifecycleManager, setupGracefulShutdown } from '../ServerLifecycle';

describe('ServerLifecycleManager', () => {
  const createServer = (): FastMCP => ({}) as FastMCP;

  it('runs start hook and tracks running state', async () => {
    const onStart = jest.fn();
    const manager = new ServerLifecycleManager(createServer(), { onStart });

    await manager.start();

    expect(onStart).toHaveBeenCalled();
    expect(manager.isServerRunning()).toBe(true);
  });

  it('prevents starting twice', async () => {
    const manager = new ServerLifecycleManager(createServer());
    await manager.start();

    await expect(manager.start()).rejects.toThrow('Server is already running');
  });

  it('invokes stop hook and resets state', async () => {
    const onStop = jest.fn();
    const manager = new ServerLifecycleManager(createServer(), { onStop });

    await manager.start();
    await manager.stop();

    expect(onStop).toHaveBeenCalled();
    expect(manager.isServerRunning()).toBe(false);
  });

  it('swallows stop when not running', async () => {
    const manager = new ServerLifecycleManager(createServer());
    await expect(manager.stop()).resolves.toBeUndefined();
  });

  it('triggers error hook when start fails', async () => {
    const error = new Error('boom');
    const onStart = jest.fn().mockRejectedValue(error);
    const onError = jest.fn();
    const manager = new ServerLifecycleManager(createServer(), { onStart, onError });

    await expect(manager.start()).rejects.toThrow('boom');
    expect(onError).toHaveBeenCalledWith(error, manager.getServer());
    expect(manager.isServerRunning()).toBe(false);
  });

  it('triggers error hook when stop fails', async () => {
    const error = new Error('fail');
    const onStop = jest.fn().mockRejectedValue(error);
    const onError = jest.fn();
    const manager = new ServerLifecycleManager(createServer(), { onStop, onError });

    await manager.start();
    await expect(manager.stop()).rejects.toThrow('fail');
    expect(onError).toHaveBeenCalledWith(error, manager.getServer());
    expect(manager.isServerRunning()).toBe(true);
  });
});

describe('setupGracefulShutdown', () => {
  let originalExit: typeof process.exit;
  let exitMock: jest.SpyInstance;
  let logSpy: jest.SpyInstance;
  let errorSpy: jest.SpyInstance;

  beforeEach(() => {
    originalExit = process.exit;
    exitMock = jest.spyOn(process, 'exit').mockImplementation((() => undefined) as any);
    logSpy = jest.spyOn(console, 'log').mockImplementation(() => {});
    errorSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
  });

  afterEach(() => {
    exitMock.mockRestore();
    logSpy.mockRestore();
    errorSpy.mockRestore();
    process.exit = originalExit;
  });

  it('registers signal handlers that stop the manager', async () => {
    const manager = new ServerLifecycleManager({} as FastMCP);
    const stopSpy = jest.spyOn(manager, 'stop').mockResolvedValue();

    const handlers: Record<string, () => Promise<void>> = {};
    const onSpy = jest.spyOn(process, 'on').mockImplementation(((signal: NodeJS.Signals, handler: () => Promise<void>) => {
      handlers[signal] = handler;
      return process;
    }) as any);

    setupGracefulShutdown(manager);

    expect(onSpy).toHaveBeenCalledTimes(3);
    expect(Object.keys(handlers)).toEqual(expect.arrayContaining(['SIGINT', 'SIGTERM', 'SIGQUIT']));

    await handlers['SIGINT']();

    expect(stopSpy).toHaveBeenCalled();
    expect(exitMock).toHaveBeenCalledWith(0);

    onSpy.mockRestore();
  });

  it('logs and exits with error when stop fails', async () => {
    const manager = new ServerLifecycleManager({} as FastMCP);
    const stopSpy = jest.spyOn(manager, 'stop').mockRejectedValue(new Error('shutdown failure'));

    let capturedHandler: (() => Promise<void>) | undefined;
    const onSpy = jest.spyOn(process, 'on').mockImplementation(((signal: NodeJS.Signals, handler: () => Promise<void>) => {
      if (signal === 'SIGTERM') {
        capturedHandler = handler;
      }
      return process;
    }) as any);

    setupGracefulShutdown(manager);
    await capturedHandler?.();

    expect(stopSpy).toHaveBeenCalled();
    expect(errorSpy).toHaveBeenCalledWith('Error during shutdown:', expect.any(Error));
    expect(exitMock).toHaveBeenCalledWith(1);

    onSpy.mockRestore();
  });
});