import { describe, it, expect, vi, afterEach } from 'vitest';
import {
  createRequestDeadline,
  DeadlineExceededError,
  TOOLCALL_DEADLINE_MS,
} from '../src/http/deadline.js';
import { TOOLCALL_TIMEOUT_MS } from '../src/http/mcpd-client.js';

afterEach(() => { vi.restoreAllMocks(); });

describe('TOOLCALL_DEADLINE_MS', () => {
  it('exceeds mcpd\'s own tool-call timeout', () => {
    // The watchdog arms earlier in the request than mcpd's fetch does. At equal
    // values it would always fire first, masking mcpd's specific
    // UpstreamTimeoutError with a generic deadline message and cutting off
    // pagination and the rest of the post-processing.
    expect(TOOLCALL_DEADLINE_MS).toBeGreaterThan(TOOLCALL_TIMEOUT_MS);
  });

  it('stays well below a typical client timeout so our error arrives first', () => {
    expect(TOOLCALL_DEADLINE_MS).toBeLessThan(600_000);
  });
});

describe('createRequestDeadline', () => {
  it('rejects with DeadlineExceededError when the budget runs out', async () => {
    const d = createRequestDeadline('tools/call', 40);
    await expect(Promise.race([
      new Promise<never>(() => { /* the hang we are protecting against */ }),
      d.expiry,
    ])).rejects.toBeInstanceOf(DeadlineExceededError);
    d.dispose();
  });

  it('does not fire once disposed', async () => {
    const d = createRequestDeadline('tools/call', 30);
    d.dispose();
    const outcome = await Promise.race([
      d.expiry.then(() => 'fired', () => 'fired'),
      new Promise((r) => setTimeout(() => { r('quiet'); }, 120)),
    ]);
    expect(outcome).toBe('quiet');
  });

  it('stops the clock while suspended, so an operator pause cannot time out', async () => {
    const d = createRequestDeadline('tools/call', 100);

    // A pause queue wait far longer than the whole deadline.
    const held = await d.suspended(async () => {
      await new Promise((r) => setTimeout(r, 250));
      return 'released';
    });
    expect(held).toBe('released');

    // Still alive afterwards, with roughly the original budget left.
    const outcome = await Promise.race([
      d.expiry.then(() => 'fired', () => 'fired'),
      new Promise((r) => setTimeout(() => { r('still-running'); }, 40)),
    ]);
    expect(outcome).toBe('still-running');
    d.dispose();
  });

  it('re-arms after suspension rather than becoming immortal', async () => {
    const d = createRequestDeadline('tools/call', 60);
    await d.suspended(async () => { await new Promise((r) => setTimeout(r, 100)); });
    // The remaining budget must still expire — suspending is not a bypass.
    await expect(Promise.race([
      new Promise<never>(() => { /* hang */ }),
      d.expiry,
    ])).rejects.toBeInstanceOf(DeadlineExceededError);
    d.dispose();
  });
});

describe('per-server deadline override', () => {
  it('falls back to the global when a server declares nothing', async () => {
    const { McpRouter } = await import('../src/router.js');
    const router = new McpRouter();
    expect(router.getToolCallDeadlineMs('docmost/search')).toBeUndefined();
  });

  it('returns the server override for any of its tools', async () => {
    const { McpRouter } = await import('../src/router.js');
    const router = new McpRouter();
    router.setServerDeadline('slowserver', 300_000);
    // Resolved from the canonical `server/tool` name, so every tool on that
    // server inherits it without being enumerated.
    expect(router.getToolCallDeadlineMs('slowserver/render')).toBe(300_000);
    expect(router.getToolCallDeadlineMs('slowserver/export')).toBe(300_000);
    expect(router.getToolCallDeadlineMs('other/tool')).toBeUndefined();
  });

  it('clears the override when the server stops declaring one', async () => {
    const { McpRouter } = await import('../src/router.js');
    const router = new McpRouter();
    router.setServerDeadline('slowserver', 300_000);
    router.setServerDeadline('slowserver', undefined);
    expect(router.getToolCallDeadlineMs('slowserver/render')).toBeUndefined();
  });
});