import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { McpRouter } from '../src/router.js';
import type { UpstreamConnection, JsonRpcRequest, JsonRpcResponse } from '../src/types.js';

function mockUpstream(name: string, opts: { tools?: Array<{ name: string }>; resources?: Array<{ uri: string }>; err?: string } = {}): UpstreamConnection {
  return {
    name,
    isAlive: () => true,
    close: async () => {},
    send: vi.fn(async (req: JsonRpcRequest): Promise<JsonRpcResponse> => {
      if (opts.err) {
        return { jsonrpc: '2.0', id: req.id, error: { code: -32603, message: opts.err } };
      }
      if (req.method === 'tools/list') {
        return { jsonrpc: '2.0', id: req.id, result: { tools: opts.tools ?? [] } };
      }
      if (req.method === 'resources/list') {
        return { jsonrpc: '2.0', id: req.id, result: { resources: opts.resources ?? [] } };
      }
      return { jsonrpc: '2.0', id: req.id, error: { code: -32601, message: 'not handled' } };
    }),
  } as UpstreamConnection;
}

describe('McpRouter discovery cache', () => {
  let router: McpRouter;

  beforeEach(() => {
    router = new McpRouter();
    vi.useFakeTimers();
    vi.setSystemTime(new Date('2026-04-15T12:00:00Z'));
  });

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

  it('serves tools/list from cache on the second call within TTL', async () => {
    const upstream = mockUpstream('slack', { tools: [{ name: 'search' }] });
    router.addUpstream(upstream);

    await router.discoverTools();
    await router.discoverTools();

    expect(upstream.send).toHaveBeenCalledTimes(1);
  });

  it('re-fetches after positive TTL expires', async () => {
    const upstream = mockUpstream('slack', { tools: [{ name: 'search' }] });
    router.addUpstream(upstream);

    await router.discoverTools();
    vi.advanceTimersByTime(31_000);
    await router.discoverTools();

    expect(upstream.send).toHaveBeenCalledTimes(2);
  });

  it('negative cache prevents repeated calls to a failing upstream', async () => {
    const upstream = mockUpstream('broken', { err: 'mcpd proxy error: timeout' });
    router.addUpstream(upstream);

    await router.discoverTools();
    await router.discoverTools();
    await router.discoverTools();

    expect(upstream.send).toHaveBeenCalledTimes(1);
  });

  it('negative cache expires after negative TTL', async () => {
    const upstream = mockUpstream('broken', { err: 'mcpd proxy error: timeout' });
    router.addUpstream(upstream);

    await router.discoverTools();
    vi.advanceTimersByTime(31_000);
    await router.discoverTools();

    expect(upstream.send).toHaveBeenCalledTimes(2);
  });

  it('re-registering a server invalidates its cache entry', async () => {
    const upstream1 = mockUpstream('slack', { tools: [{ name: 'v1' }] });
    router.addUpstream(upstream1);
    await router.discoverTools();
    expect(upstream1.send).toHaveBeenCalledTimes(1);

    const upstream2 = mockUpstream('slack', { tools: [{ name: 'v2' }] });
    router.addUpstream(upstream2);
    const tools = await router.discoverTools();

    expect(upstream2.send).toHaveBeenCalledTimes(1);
    expect(tools.map((t) => t.name)).toEqual(['slack/v2']);
  });

  it('removeUpstream clears cache so follow-up add re-fetches', async () => {
    const upstream1 = mockUpstream('slack', { tools: [{ name: 'v1' }] });
    router.addUpstream(upstream1);
    await router.discoverTools();

    router.removeUpstream('slack');

    const upstream2 = mockUpstream('slack', { tools: [{ name: 'v2' }] });
    router.addUpstream(upstream2);
    await router.discoverTools();

    expect(upstream2.send).toHaveBeenCalledTimes(1);
  });

  it('one dead server does not block cached results for others', async () => {
    const broken = mockUpstream('broken', { err: 'timeout' });
    const healthy = mockUpstream('healthy', { tools: [{ name: 'ping' }] });
    router.addUpstream(broken);
    router.addUpstream(healthy);

    const first = await router.discoverTools();
    expect(first.map((t) => t.name)).toEqual(['healthy/ping']);

    // Second call: both come from cache.
    const second = await router.discoverTools();
    expect(second.map((t) => t.name)).toEqual(['healthy/ping']);
    expect(broken.send).toHaveBeenCalledTimes(1);
    expect(healthy.send).toHaveBeenCalledTimes(1);
  });

  it('discoverResources uses its own cache key independent of tools/list', async () => {
    const upstream = mockUpstream('docs', { tools: [{ name: 'search' }], resources: [{ uri: 'doc://1' }] });
    router.addUpstream(upstream);

    await router.discoverTools();
    await router.discoverResources();
    await router.discoverTools();
    await router.discoverResources();

    // Each method cached separately → exactly one call per method.
    expect(upstream.send).toHaveBeenCalledTimes(2);
  });
});