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

function mockUpstream(
  name: string,
  opts: {
    tools?: Array<{ name: string; description?: string }>;
    resources?: Array<{ uri: string; name?: string; description?: string }>;
    prompts?: Array<{ name: string; description?: string }>;
  } = {},
): UpstreamConnection {
  const notificationHandlers: Array<(n: JsonRpcNotification) => void> = [];
  return {
    name,
    isAlive: vi.fn(() => true),
    close: vi.fn(async () => {}),
    onNotification: vi.fn((handler: (n: JsonRpcNotification) => void) => {
      notificationHandlers.push(handler);
    }),
    send: vi.fn(async (req: JsonRpcRequest): Promise<JsonRpcResponse> => {
      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 ?? [] },
        };
      }
      if (req.method === 'prompts/list') {
        return {
          jsonrpc: '2.0',
          id: req.id,
          result: { prompts: opts.prompts ?? [] },
        };
      }
      if (req.method === 'tools/call') {
        return {
          jsonrpc: '2.0',
          id: req.id,
          result: {
            content: [{ type: 'text', text: `Called ${(req.params as Record<string, unknown>)?.name}` }],
          },
        };
      }
      if (req.method === 'resources/read') {
        return {
          jsonrpc: '2.0',
          id: req.id,
          result: {
            contents: [{ uri: (req.params as Record<string, unknown>)?.uri, text: 'resource content' }],
          },
        };
      }
      if (req.method === 'prompts/get') {
        return {
          jsonrpc: '2.0',
          id: req.id,
          result: {
            messages: [{ role: 'user', content: { type: 'text', text: 'prompt content' } }],
          },
        };
      }
      return { jsonrpc: '2.0', id: req.id, error: { code: -32601, message: 'Not found' } };
    }),
    // expose for tests
    _notificationHandlers: notificationHandlers,
  } as UpstreamConnection & { _notificationHandlers: Array<(n: JsonRpcNotification) => void> };
}

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

  beforeEach(() => {
    router = new McpRouter();
  });

  describe('initialize', () => {
    it('responds with server info and capabilities including resources and prompts', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'initialize',
      });

      expect(res.result).toBeDefined();
      const result = res.result as Record<string, unknown>;
      expect(result['protocolVersion']).toBe('2024-11-05');
      expect((result['serverInfo'] as Record<string, unknown>)['name']).toBe('mcpctl-proxy');
      const capabilities = result['capabilities'] as Record<string, unknown>;
      expect(capabilities['tools']).toBeDefined();
      expect(capabilities['resources']).toBeDefined();
      expect(capabilities['prompts']).toBeDefined();
    });
  });

  describe('tools/list', () => {
    it('returns empty tools when no upstreams', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'tools/list',
      });

      const result = res.result as { tools: unknown[] };
      expect(result.tools).toEqual([]);
    });

    it('discovers and namespaces tools from upstreams', async () => {
      router.addUpstream(mockUpstream('slack', {
        tools: [
          { name: 'send_message', description: 'Send a message' },
          { name: 'list_channels', description: 'List channels' },
        ],
      }));
      router.addUpstream(mockUpstream('github', {
        tools: [
          { name: 'create_issue', description: 'Create an issue' },
        ],
      }));

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'tools/list',
      });

      const result = res.result as { tools: Array<{ name: string }> };
      expect(result.tools).toHaveLength(3);
      expect(result.tools.map((t) => t.name)).toContain('slack/send_message');
      expect(result.tools.map((t) => t.name)).toContain('slack/list_channels');
      expect(result.tools.map((t) => t.name)).toContain('github/create_issue');
    });

    it('skips unavailable upstreams', async () => {
      const failingUpstream = mockUpstream('failing');
      vi.mocked(failingUpstream.send).mockRejectedValue(new Error('Connection refused'));
      router.addUpstream(failingUpstream);

      router.addUpstream(mockUpstream('working', {
        tools: [{ name: 'do_thing', description: 'Does a thing' }],
      }));

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'tools/list',
      });

      const result = res.result as { tools: Array<{ name: string }> };
      expect(result.tools).toHaveLength(1);
      expect(result.tools[0]?.name).toBe('working/do_thing');
    });
  });

  describe('tools/call', () => {
    it('routes call to correct upstream', async () => {
      const slack = mockUpstream('slack', { tools: [{ name: 'send_message' }] });
      router.addUpstream(slack);
      await router.discoverTools();

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'tools/call',
        params: { name: 'slack/send_message', arguments: { channel: '#general', text: 'hello' } },
      });

      expect(res.result).toBeDefined();
      // Verify the upstream received the call with de-namespaced tool name
      expect(vi.mocked(slack.send)).toHaveBeenCalledWith(
        expect.objectContaining({
          method: 'tools/call',
          params: expect.objectContaining({ name: 'send_message' }),
        }),
      );
    });

    it('returns error for unknown tool', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'tools/call',
        params: { name: 'unknown/tool' },
      });

      expect(res.error).toBeDefined();
      expect(res.error?.code).toBe(-32601);
    });

    it('returns error when tool name is missing', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'tools/call',
        params: {},
      });

      expect(res.error).toBeDefined();
      expect(res.error?.code).toBe(-32602);
    });

    it('returns error when upstream is dead', async () => {
      const slack = mockUpstream('slack', { tools: [{ name: 'send_message' }] });
      router.addUpstream(slack);
      await router.discoverTools();

      vi.mocked(slack.isAlive).mockReturnValue(false);

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'tools/call',
        params: { name: 'slack/send_message' },
      });

      expect(res.error).toBeDefined();
      expect(res.error?.code).toBe(-32603);
    });
  });

  describe('resources/list', () => {
    it('returns empty resources when no upstreams', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'resources/list',
      });

      const result = res.result as { resources: unknown[] };
      expect(result.resources).toEqual([]);
    });

    it('discovers and namespaces resources from upstreams', async () => {
      router.addUpstream(mockUpstream('files', {
        resources: [
          { uri: 'file:///docs/readme.md', name: 'README', description: 'Project readme' },
        ],
      }));
      router.addUpstream(mockUpstream('db', {
        resources: [
          { uri: 'db://users', name: 'Users table' },
        ],
      }));

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'resources/list',
      });

      const result = res.result as { resources: Array<{ uri: string }> };
      expect(result.resources).toHaveLength(2);
      expect(result.resources.map((r) => r.uri)).toContain('files://file:///docs/readme.md');
      expect(result.resources.map((r) => r.uri)).toContain('db://db://users');
    });
  });

  describe('resources/read', () => {
    it('routes read to correct upstream', async () => {
      const files = mockUpstream('files', {
        resources: [{ uri: 'file:///docs/readme.md', name: 'README' }],
      });
      router.addUpstream(files);
      await router.discoverResources();

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'resources/read',
        params: { uri: 'files://file:///docs/readme.md' },
      });

      expect(res.result).toBeDefined();
      expect(vi.mocked(files.send)).toHaveBeenCalledWith(
        expect.objectContaining({
          method: 'resources/read',
          params: expect.objectContaining({ uri: 'file:///docs/readme.md' }),
        }),
      );
    });

    it('returns error for unknown resource', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'resources/read',
        params: { uri: 'unknown://resource' },
      });

      expect(res.error).toBeDefined();
      expect(res.error?.code).toBe(-32601);
    });
  });

  describe('prompts/list', () => {
    it('returns empty prompts when no upstreams', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'prompts/list',
      });

      const result = res.result as { prompts: unknown[] };
      expect(result.prompts).toEqual([]);
    });

    it('discovers and namespaces prompts from upstreams', async () => {
      router.addUpstream(mockUpstream('code', {
        prompts: [
          { name: 'review', description: 'Code review' },
          { name: 'explain', description: 'Explain code' },
        ],
      }));

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'prompts/list',
      });

      const result = res.result as { prompts: Array<{ name: string }> };
      expect(result.prompts).toHaveLength(2);
      expect(result.prompts.map((p) => p.name)).toContain('code/review');
      expect(result.prompts.map((p) => p.name)).toContain('code/explain');
    });
  });

  describe('prompts/get', () => {
    it('routes get to correct upstream', async () => {
      const code = mockUpstream('code', {
        prompts: [{ name: 'review', description: 'Code review' }],
      });
      router.addUpstream(code);
      await router.discoverPrompts();

      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'prompts/get',
        params: { name: 'code/review' },
      });

      expect(res.result).toBeDefined();
      expect(vi.mocked(code.send)).toHaveBeenCalledWith(
        expect.objectContaining({
          method: 'prompts/get',
          params: expect.objectContaining({ name: 'review' }),
        }),
      );
    });

    it('returns error for unknown prompt', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'prompts/get',
        params: { name: 'unknown/prompt' },
      });

      expect(res.error).toBeDefined();
      expect(res.error?.code).toBe(-32601);
    });
  });

  describe('notifications', () => {
    it('forwards notifications from upstreams with source tag', () => {
      const received: JsonRpcNotification[] = [];
      router.setNotificationHandler((n) => received.push(n));

      const slack = mockUpstream('slack', { tools: [{ name: 'send_message' }] }) as UpstreamConnection & {
        _notificationHandlers: Array<(n: JsonRpcNotification) => void>;
      };
      router.addUpstream(slack);

      // Simulate upstream sending a notification
      for (const handler of slack._notificationHandlers) {
        handler({ jsonrpc: '2.0', method: 'notifications/progress', params: { progress: 50 } });
      }

      expect(received).toHaveLength(1);
      expect(received[0]?.method).toBe('notifications/progress');
      expect(received[0]?.params?._source).toBe('slack');
    });
  });

  describe('unknown methods', () => {
    it('returns method not found error', async () => {
      const res = await router.route({
        jsonrpc: '2.0',
        id: 1,
        method: 'completions/complete',
      });

      expect(res.error).toBeDefined();
      expect(res.error?.code).toBe(-32601);
    });
  });

  describe('upstream management', () => {
    it('lists upstream names', () => {
      router.addUpstream(mockUpstream('slack'));
      router.addUpstream(mockUpstream('github'));
      expect(router.getUpstreamNames()).toEqual(['slack', 'github']);
    });

    it('removes upstream and cleans up all mappings', async () => {
      const slack = mockUpstream('slack', {
        tools: [{ name: 'send_message' }],
        resources: [{ uri: 'slack://channels' }],
        prompts: [{ name: 'compose' }],
      });
      router.addUpstream(slack);
      await router.discoverTools();
      await router.discoverResources();
      await router.discoverPrompts();

      router.removeUpstream('slack');
      expect(router.getUpstreamNames()).toEqual([]);

      // Verify tool/resource/prompt mappings are cleaned
      const toolRes = await router.route({
        jsonrpc: '2.0', id: 1, method: 'tools/call',
        params: { name: 'slack/send_message' },
      });
      expect(toolRes.error?.code).toBe(-32601);
    });

    it('closes all upstreams', async () => {
      const slack = mockUpstream('slack');
      const github = mockUpstream('github');
      router.addUpstream(slack);
      router.addUpstream(github);

      await router.closeAll();

      expect(slack.close).toHaveBeenCalled();
      expect(github.close).toHaveBeenCalled();
      expect(router.getUpstreamNames()).toEqual([]);
    });
  });

  describe('onUpstreamCall for discovery methods', () => {
    let onUpstreamCall: ReturnType<typeof vi.fn>;

    beforeEach(() => {
      onUpstreamCall = vi.fn();
      router.onUpstreamCall = onUpstreamCall;
      router.addUpstream(mockUpstream('slack', {
        tools: [{ name: 'send_message' }],
        resources: [{ uri: 'slack://channels' }],
        prompts: [{ name: 'compose' }],
      }));
      router.addUpstream(mockUpstream('ha', {
        tools: [{ name: 'get_entities' }],
        resources: [{ uri: 'ha://states' }],
        prompts: [{ name: 'assist' }],
      }));
    });

    it('fires onUpstreamCall for each server during tools/list', async () => {
      await router.route({ jsonrpc: '2.0', id: 1, method: 'tools/list' });

      expect(onUpstreamCall).toHaveBeenCalledTimes(2);
      const calls = onUpstreamCall.mock.calls.map((c: unknown[]) => c[0] as { upstream: string; method: string; request: unknown; response: unknown; durationMs: number });
      expect(calls[0]!.upstream).toBe('slack');
      expect(calls[0]!.method).toBe('tools/list');
      expect(calls[0]!.durationMs).toBeGreaterThanOrEqual(0);
      expect(calls[0]!.request).toBeDefined();
      expect(calls[0]!.response).toBeDefined();
      expect(calls[1]!.upstream).toBe('ha');
      expect(calls[1]!.method).toBe('tools/list');
    });

    it('fires onUpstreamCall for each server during resources/list', async () => {
      await router.route({ jsonrpc: '2.0', id: 1, method: 'resources/list' });

      expect(onUpstreamCall).toHaveBeenCalledTimes(2);
      const calls = onUpstreamCall.mock.calls.map((c: unknown[]) => c[0] as { upstream: string; method: string });
      expect(calls[0]!.upstream).toBe('slack');
      expect(calls[0]!.method).toBe('resources/list');
      expect(calls[1]!.upstream).toBe('ha');
      expect(calls[1]!.method).toBe('resources/list');
    });

    it('fires onUpstreamCall for each server during prompts/list', async () => {
      await router.route({ jsonrpc: '2.0', id: 1, method: 'prompts/list' });

      expect(onUpstreamCall).toHaveBeenCalledTimes(2);
      const calls = onUpstreamCall.mock.calls.map((c: unknown[]) => c[0] as { upstream: string; method: string });
      expect(calls[0]!.upstream).toBe('slack');
      expect(calls[0]!.method).toBe('prompts/list');
      expect(calls[1]!.upstream).toBe('ha');
      expect(calls[1]!.method).toBe('prompts/list');
    });

    it('skips failed upstream but fires for successful ones', async () => {
      const failing = mockUpstream('failing');
      vi.mocked(failing.send).mockRejectedValue(new Error('Connection refused'));
      router.addUpstream(failing);

      await router.route({ jsonrpc: '2.0', id: 1, method: 'tools/list' });

      // slack + ha succeed, failing throws — onUpstreamCall fires only for successful ones
      expect(onUpstreamCall).toHaveBeenCalledTimes(2);
      const upstreams = onUpstreamCall.mock.calls.map((c: unknown[]) => (c[0] as { upstream: string }).upstream);
      expect(upstreams).toContain('slack');
      expect(upstreams).toContain('ha');
      expect(upstreams).not.toContain('failing');
    });

    it('does not fire onUpstreamCall when callback is null', async () => {
      router.onUpstreamCall = null;
      // Should not throw
      await router.route({ jsonrpc: '2.0', id: 1, method: 'tools/list' });
      expect(onUpstreamCall).not.toHaveBeenCalled();
    });
  });

  describe('per-server proxymodel resolution', () => {
    const mockLlm = { complete: async () => '', available: () => false };
    const mockCache = {
      getOrCompute: async (_k: string, fn: () => Promise<string>) => fn(),
      hash: () => '',
      get: async () => null,
      set: async () => {},
    };
    const haLlm = { complete: async () => 'ha-result', available: () => true };
    const haCache = {
      getOrCompute: async (_k: string, fn: () => Promise<string>) => fn(),
      hash: (s: string) => s.slice(0, 4),
      get: async () => null,
      set: async () => {},
    };

    it('uses server-specific proxymodel when set', () => {
      router.setProxyModel('default', mockLlm, mockCache);
      router.setServerProxyModel('ha', 'ha-special', haLlm, haCache);

      // Access private method via cast
      const r = router as unknown as { getProxyModelForServer(s: string): { name: string } | null };
      const config = r.getProxyModelForServer('ha');
      expect(config).not.toBeNull();
      expect(config!.name).toBe('ha-special');
      expect(config!).toHaveProperty('llm', haLlm);
      expect(config!).toHaveProperty('cache', haCache);
    });

    it('falls back to default when no server override', () => {
      router.setProxyModel('default', mockLlm, mockCache);
      router.setServerProxyModel('ha', 'ha-special', haLlm, haCache);

      const r = router as unknown as { getProxyModelForServer(s: string): { name: string } | null };
      const config = r.getProxyModelForServer('slack');
      expect(config).not.toBeNull();
      expect(config!.name).toBe('default');
      expect(config!).toHaveProperty('llm', mockLlm);
    });

    it('returns null when no default and no server override', () => {
      const r = router as unknown as { getProxyModelForServer(s: string): { name: string } | null };
      const config = r.getProxyModelForServer('slack');
      expect(config).toBeNull();
    });

    it('setServerProxyModel overwrites previous setting', () => {
      router.setServerProxyModel('ha', 'model-a', mockLlm, mockCache);
      router.setServerProxyModel('ha', 'model-b', haLlm, haCache);

      const r = router as unknown as { getProxyModelForServer(s: string): { name: string } | null };
      const config = r.getProxyModelForServer('ha');
      expect(config).not.toBeNull();
      expect(config!.name).toBe('model-b');
      expect(config!).toHaveProperty('llm', haLlm);
    });
  });

  describe('tool_call_trace audit emission', () => {
    it('emits tool_call_trace on successful tool call', async () => {
      const alpha = mockUpstream('alpha', { tools: [{ name: 'do_thing' }] });
      router.addUpstream(alpha);
      await router.discoverTools();

      const emitted: Array<Record<string, unknown>> = [];
      const mockCollector = { emit: vi.fn((e: Record<string, unknown>) => emitted.push(e)) };
      router.setAuditCollector(mockCollector as never);

      await router.route(
        { jsonrpc: '2.0', id: 1, method: 'tools/call', params: { name: 'alpha/do_thing', arguments: { key: 'val' } } },
        { sessionId: 'sess-1' },
      );

      expect(mockCollector.emit).toHaveBeenCalledOnce();
      const event = emitted[0]!;
      expect(event['eventKind']).toBe('tool_call_trace');
      expect(event['sessionId']).toBe('sess-1');
      expect(event['serverName']).toBe('alpha');
      expect(event['verified']).toBe(true);
      const payload = event['payload'] as Record<string, unknown>;
      expect(payload['toolName']).toBe('alpha/do_thing');
      expect(payload['argKeys']).toBe('key');
      expect(payload['durationMs']).toBeTypeOf('number');
      expect(payload['resultSizeBytes']).toBeTypeOf('number');
      expect(payload['error']).toBeNull();
    });

    it('does not emit when auditCollector is not set', async () => {
      const alpha = mockUpstream('alpha', { tools: [{ name: 'do_thing' }] });
      router.addUpstream(alpha);
      await router.discoverTools();

      // No setAuditCollector call — should not throw
      const resp = await router.route(
        { jsonrpc: '2.0', id: 1, method: 'tools/call', params: { name: 'alpha/do_thing', arguments: {} } },
        { sessionId: 'sess-1' },
      );
      expect(resp.result).toBeDefined();
    });

    it('does not emit when sessionId is missing', async () => {
      const alpha = mockUpstream('alpha', { tools: [{ name: 'do_thing' }] });
      router.addUpstream(alpha);
      await router.discoverTools();

      const mockCollector = { emit: vi.fn() };
      router.setAuditCollector(mockCollector as never);

      await router.route(
        { jsonrpc: '2.0', id: 1, method: 'tools/call', params: { name: 'alpha/do_thing', arguments: {} } },
      );

      expect(mockCollector.emit).not.toHaveBeenCalled();
    });

    it('captures error in trace when upstream returns error', async () => {
      const failing: UpstreamConnection = {
        name: 'fail-srv',
        isAlive: vi.fn(() => true),
        close: vi.fn(async () => {}),
        onNotification: vi.fn(),
        send: vi.fn(async (req: JsonRpcRequest): Promise<JsonRpcResponse> => {
          if (req.method === 'tools/list') {
            return { jsonrpc: '2.0', id: req.id, result: { tools: [{ name: 'fail_tool' }] } };
          }
          return { jsonrpc: '2.0', id: req.id, error: { code: -32000, message: 'Something broke' } };
        }),
      };
      router.addUpstream(failing);
      await router.discoverTools();

      const emitted: Array<Record<string, unknown>> = [];
      const mockCollector = { emit: vi.fn((e: Record<string, unknown>) => emitted.push(e)) };
      router.setAuditCollector(mockCollector as never);

      await router.route(
        { jsonrpc: '2.0', id: 1, method: 'tools/call', params: { name: 'fail-srv/fail_tool', arguments: { a: 1, b: 2 } } },
        { sessionId: 'sess-err' },
      );

      expect(mockCollector.emit).toHaveBeenCalledOnce();
      const payload = emitted[0]!['payload'] as Record<string, unknown>;
      expect(payload['error']).toBe('Something broke');
      expect(payload['argKeys']).toBe('a, b');
    });
  });
});
