import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { EventEmitter } from 'events';
import { MCPToolRegistry } from './tool-registry';
import type { MCPServerManager } from './server-manager';
import type { MCPConfig } from '@lace/agent/config/mcp-types';

// Minimal interface for test tool mocks
interface MinimalTool {
  name?: string;
  description: string;
  schema: any;
  execute: any;
}

// Typed fake implementation of MCPServerManager
class FakeMCPServerManager extends EventEmitter implements MCPServerManager {
  startServer = vi.fn();
  stopServer = vi.fn();
  getServer = vi.fn();
  getClient = vi.fn();
  getAllServers = vi.fn();
  shutdown = vi.fn();
  isServerRunning = vi.fn();
  restartServer = vi.fn();
}

describe('MCPToolRegistry', () => {
  let registry: MCPToolRegistry;
  let mockServerManager: FakeMCPServerManager;

  beforeEach(() => {
    mockServerManager = new FakeMCPServerManager();

    // Set up mock implementations
    mockServerManager.startServer.mockResolvedValue();
    mockServerManager.getClient.mockReturnValue({
      listTools: vi.fn().mockResolvedValue({
        tools: [
          {
            name: 'read_file',
            description: 'Read a file',
            inputSchema: {
              type: 'object',
              properties: { path: { type: 'string' } },
              required: ['path'],
            },
          },
        ],
      }),
    });
    mockServerManager.getServer.mockImplementation((serverId: string) => ({
      id: serverId,
      connectionKey: serverId,
      config: { command: 'node', enabled: true, tools: {} },
      status: 'running',
      client: mockServerManager.getClient(serverId),
    }));
    mockServerManager.getAllServers.mockReturnValue([]);
    mockServerManager.shutdown.mockResolvedValue();

    registry = new MCPToolRegistry(mockServerManager);
  });

  afterEach(async () => {
    await registry.shutdown();
  });

  it('should initialize and start enabled servers', async () => {
    const config: MCPConfig = {
      servers: {
        filesystem: {
          command: 'node',
          args: ['fs.js'],
          enabled: true,
          tools: { read_file: 'allow' },
        },
        browser: {
          command: 'python',
          args: ['browser.py'],
          enabled: false, // Should not start this one
          tools: {},
        },
      },
    };

    await registry.initialize(config);

    expect(mockServerManager.startServer).toHaveBeenCalledWith(
      expect.objectContaining({
        serverId: 'filesystem',
        config: { ...config.servers.filesystem, placement: 'host' },
        runtime: expect.objectContaining({ id: 'mcp-registry:filesystem' }),
      })
    );
    expect(mockServerManager.startServer).not.toHaveBeenCalledWith(
      expect.objectContaining({ serverId: 'browser' })
    );
  });

  it('should discover tools when server comes online', async () => {
    // Create promise that resolves when tools-updated event is emitted
    const toolsUpdatedPromise = new Promise<[string, unknown[]]>((resolve) => {
      registry.once('tools-updated', (serverId, tools) => {
        resolve([serverId, tools]);
      });
    });

    // Simulate server coming online
    mockServerManager.emit('server-status-changed', 'filesystem', 'running');

    // Wait for the actual event instead of arbitrary timeout
    const [serverId, tools] = await toolsUpdatedPromise;

    expect(serverId).toBe('filesystem');
    expect(tools).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          name: 'filesystem/read_file',
          description: 'Read a file',
        }),
      ])
    );
  });

  it('should filter disabled tools from available tools', () => {
    // Manually set up tools for testing
    const mockTool: MinimalTool = {
      description: 'Write file',
      schema: vi.fn(),
      execute: vi.fn(),
    };

    registry['toolsByConnection'].set('filesystem', [
      { name: 'filesystem/read_file', ...mockTool },
      { name: 'filesystem/write_file', ...mockTool },
    ]);

    const config: MCPConfig = {
      servers: {
        filesystem: {
          command: 'node',
          enabled: true,
          tools: {
            read_file: 'allow',
            write_file: 'disable', // This should be filtered out
          },
        },
      },
    };

    const availableTools = registry.getAvailableTools(config);

    expect(availableTools).toHaveLength(1);
    expect(availableTools[0].name).toBe('filesystem/read_file');
  });

  it('should get correct approval level for tools', () => {
    const config: MCPConfig = {
      servers: {
        filesystem: {
          command: 'node',
          enabled: true,
          tools: {
            read_file: 'allow',
            write_file: 'ask',
          },
        },
      },
    };

    expect(registry.getToolApprovalLevel(config, 'filesystem/read_file')).toBe('allow');
    expect(registry.getToolApprovalLevel(config, 'filesystem/write_file')).toBe('ask');
    expect(registry.getToolApprovalLevel(config, 'filesystem/unknown_tool')).toBe('ask');
  });

  it('should handle tool discovery errors gracefully', async () => {
    const errorHandler = vi.fn();
    registry.on('tool-discovery-error', errorHandler);

    // Mock getClient to return a client that throws an error
    vi.spyOn(mockServerManager, 'getClient').mockReturnValue({
      listTools: vi.fn().mockRejectedValue(new Error('Server unavailable')),
    } as any);

    // Simulate server coming online
    mockServerManager.emit('server-status-changed', 'filesystem', 'running');

    // Wait for async error handling
    await new Promise((resolve) => setTimeout(resolve, 10));

    expect(errorHandler).toHaveBeenCalledWith('filesystem', 'Server unavailable');
  });

  it('should clear tools when server stops', () => {
    // Set up tools first
    const mockTool = { name: 'filesystem/read_file' } as any;
    registry['toolsByConnection'].set('filesystem', [mockTool]);

    const toolsUpdated = vi.fn();
    registry.on('tools-updated', toolsUpdated);

    // Simulate server stopping
    mockServerManager.emit('server-status-changed', 'filesystem', 'stopped');

    expect(registry.getServerTools('filesystem')).toHaveLength(0);
    expect(toolsUpdated).toHaveBeenCalledWith('filesystem', []);
  });
});
