import { describe, it, expect, beforeEach, vi } from 'vitest';
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { setupRAGRetrieval } from '../index.js';
import { ConfigManager } from '../../../config/config-manager.js';
import { promises as fs } from 'fs';

vi.mock('fs', () => ({
  promises: {
    readFile: vi.fn(),
    readdir: vi.fn(),
    stat: vi.fn(),
    mkdir: vi.fn(),
    writeFile: vi.fn(),
    access: vi.fn()
  }
}));

const mockFs = vi.mocked(fs);

describe('RAG Retrieval MCP Integration', () => {
  let server: Server;
  let configManager: ConfigManager;
  let mockSetRequestHandler: ReturnType<typeof vi.fn>;

  beforeEach(() => {
    vi.clearAllMocks();
    
    mockSetRequestHandler = vi.fn();
    server = {
      setRequestHandler: mockSetRequestHandler
    } as any;

    const mockConfig = {
      rag: {
        enabled: true,
        indexPaths: ['./docs'],
        chunkSize: 500,
        chunkOverlap: 50,
        embeddingModel: 'test-model',
        collections: {
          docs: {
            name: 'docs',
            paths: ['./docs'],
            description: 'Documentation'
          }
        },
        persistencePath: '.atlas/rag',
        autoIndex: false,
        cacheEmbeddings: true
      }
    };
    
    configManager = {
      getProjectPath: vi.fn().mockReturnValue('/project'),
      config: mockConfig,
      get: vi.fn().mockReturnValue(mockConfig)
    } as any;
  });

  describe('setup', () => {
    it('should return tool definitions when enabled', async () => {
      const result = await setupRAGRetrieval(server, configManager);

      expect(result).toBeDefined();
      expect(result.tools).toBeDefined();
      expect(result.tools.length).toBeGreaterThan(0);
    });

    it('should return empty tools when disabled', async () => {
      configManager.config.rag.enabled = false;
      (configManager.get as any).mockReturnValue({
        ...configManager.config,
        rag: { ...configManager.config.rag, enabled: false }
      });

      const result = await setupRAGRetrieval(server, configManager);

      expect(result.tools).toHaveLength(0);
    });

    it('should register all RAG tools', async () => {
      const result = await setupRAGRetrieval(server, configManager);

      const toolNames = result.tools.map((t: any) => t.name);
      expect(toolNames).toContain('rag_search');
      expect(toolNames).toContain('rag_index_document');
      expect(toolNames).toContain('rag_index_directory');
      expect(toolNames).toContain('rag_index_collection');
      expect(toolNames).toContain('rag_get_stats');
      expect(toolNames).toContain('rag_clear_index');
    });

    it('should register request handler', async () => {
      await setupRAGRetrieval(server, configManager);

      expect(mockSetRequestHandler).toHaveBeenCalled();
    });
  });

  describe('tool handlers', () => {
    let handler: any;

    beforeEach(async () => {
      await setupRAGRetrieval(server, configManager);
      handler = mockSetRequestHandler.mock.calls[0][1];
    });

    describe('rag_search', () => {
      it('should search documents', async () => {
        // Initialize the pipeline first
        mockFs.access.mockRejectedValueOnce(new Error('Not found'));
        mockFs.mkdir.mockResolvedValueOnce(undefined);
        await setupRAGRetrieval(server, configManager);
        
        // Index a document first
        mockFs.readFile.mockResolvedValueOnce('Test content about machine learning');
        await handler({
          method: 'tools/call',
          params: {
            name: 'rag_index_document',
            arguments: { path: './test.md' }
          }
        });

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_search',
            arguments: {
              query: 'machine learning',
              limit: 5
            }
          }
        };

        const response = await handler(request);

        expect(response.content).toBeDefined();
        expect(response.content[0].type).toBe('text');
        expect(response.content[0].text).toContain('Found');
      });

      it('should handle search with filters', async () => {
        mockFs.access.mockResolvedValueOnce(undefined);
        mockFs.readFile.mockResolvedValueOnce(JSON.stringify({
          chunks: [],
          documentChunks: []
        }));

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_search',
            arguments: {
              query: 'test',
              limit: 5,
              filters: { language: 'python' }
            }
          }
        };

        const response = await handler(request);

        expect(response.content).toBeDefined();
        expect(response.content[0].text).toContain('No results found');
      });
    });

    describe('rag_index_document', () => {
      it('should index a document', async () => {
        mockFs.access.mockRejectedValueOnce(new Error('Not found'));
        mockFs.mkdir.mockResolvedValueOnce(undefined);
        mockFs.readFile.mockResolvedValueOnce('# Test Document\n\nContent here.');
        mockFs.writeFile.mockResolvedValueOnce(undefined);

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_index_document',
            arguments: {
              path: './docs/test.md'
            }
          }
        };

        const response = await handler(request);

        expect(response.content[0].text).toContain('Successfully indexed');
        expect(mockFs.readFile).toHaveBeenCalledWith(
          expect.stringContaining('test.md'),
          'utf-8'
        );
      });

      it('should handle indexing errors', async () => {
        mockFs.access.mockRejectedValueOnce(new Error('Not found'));
        mockFs.mkdir.mockResolvedValueOnce(undefined);
        mockFs.readFile.mockRejectedValueOnce(new Error('File not found'));

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_index_document',
            arguments: {
              path: './docs/missing.md'
            }
          }
        };

        const response = await handler(request);

        expect(response.content[0].text).toContain('Failed to index');
        expect(response.content[0].text).toContain('File not found');
      });
    });

    describe('rag_index_directory', () => {
      it('should index directory', async () => {
        mockFs.access.mockRejectedValueOnce(new Error('Not found'));
        mockFs.mkdir.mockResolvedValueOnce(undefined);
        mockFs.readdir.mockResolvedValueOnce(['file1.md', 'file2.md'] as any);
        mockFs.stat.mockResolvedValue({ isFile: () => true, isDirectory: () => false } as any);
        mockFs.readFile
          .mockResolvedValueOnce('# File 1')
          .mockResolvedValueOnce('# File 2')
          .mockResolvedValueOnce(JSON.stringify({ chunks: [], documentChunks: [] }));
        mockFs.writeFile.mockResolvedValue(undefined);

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_index_directory',
            arguments: {
              path: './docs'
            }
          }
        };

        const response = await handler(request);

        expect(response.content[0].text).toContain('Successfully indexed 2 documents');
      });
    });

    describe('rag_index_collection', () => {
      it('should index collection', async () => {
        mockFs.access.mockRejectedValueOnce(new Error('Not found'));
        mockFs.mkdir.mockResolvedValueOnce(undefined);
        mockFs.readdir.mockResolvedValueOnce(['test.md'] as any);
        mockFs.stat.mockResolvedValue({ isFile: () => true, isDirectory: () => false } as any);
        mockFs.readFile
          .mockResolvedValueOnce('# Test')
          .mockResolvedValueOnce(JSON.stringify({ chunks: [], documentChunks: [] }));
        mockFs.writeFile.mockResolvedValue(undefined);

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_index_collection',
            arguments: {
              collection: 'docs'
            }
          }
        };

        const response = await handler(request);

        expect(response.content[0].text).toContain('Successfully indexed collection');
      });

      it('should handle non-existent collection', async () => {
        mockFs.access.mockRejectedValueOnce(new Error('Not found'));
        mockFs.mkdir.mockResolvedValueOnce(undefined);

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_index_collection',
            arguments: {
              collection: 'non-existent'
            }
          }
        };

        const response = await handler(request);

        expect(response.content[0].text).toContain('Failed to index');
        expect(response.content[0].text).toContain('Collection not found');
      });
    });

    describe('rag_get_stats', () => {
      it('should return statistics', async () => {
        // Initialize the pipeline first
        mockFs.access.mockRejectedValueOnce(new Error('Not found'));
        mockFs.mkdir.mockResolvedValueOnce(undefined);
        await setupRAGRetrieval(server, configManager);
        
        // Index a document to have some stats
        mockFs.readFile.mockResolvedValueOnce('Test content');
        await handler({
          method: 'tools/call',
          params: {
            name: 'rag_index_document',
            arguments: { path: './test.md' }
          }
        });

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_get_stats',
            arguments: {}
          }
        };

        const response = await handler(request);

        expect(response.content[0].text).toContain('RAG Statistics');
        expect(response.content[0].text).toContain('Total Documents: 1');
        expect(response.content[0].text).toContain('Total Chunks: 1');
      });
    });

    describe('rag_clear_index', () => {
      it('should clear index', async () => {
        mockFs.access.mockResolvedValueOnce(undefined);
        mockFs.readFile.mockResolvedValueOnce(JSON.stringify({
          chunks: [{ id: 'chunk-1', chunk: {}, embedding: [] }],
          documentChunks: []
        }));
        mockFs.writeFile.mockResolvedValueOnce(undefined);

        const request = {
          method: 'tools/call',
          params: {
            name: 'rag_clear_index',
            arguments: {}
          }
        };

        const response = await handler(request);

        expect(response.content[0].text).toContain('Successfully cleared');
        expect(mockFs.writeFile).toHaveBeenCalled();
      });
    });
  });
});