import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { RAGPipeline } from './rag-pipeline.js';
import { LocalEmbeddingModel } from './embeddings.js';
import { InMemoryVectorStore } from './vector-store.js';
import { MarkdownDocumentProcessor } from './document-processor.js';
import path from 'path';
let ragPipeline = null;
export async function setupRAGRetrieval(server, configManager) {
    // Check if RAG is explicitly disabled in config
    const config = configManager.get();
    if (config.rag?.enabled === false) {
        return { tools: [] };
    }
    // Default RAG configuration
    const ragConfig = {
        enabled: true,
        indexPaths: ['./docs', './README.md'],
        chunkSize: 500,
        chunkOverlap: 50,
        embeddingModel: 'all-MiniLM-L6-v2',
        collections: {
            docs: {
                name: 'docs',
                paths: ['./docs'],
                description: 'Documentation files'
            },
            readme: {
                name: 'readme',
                paths: ['./README.md', './docs/README.md'],
                description: 'README files'
            }
        },
        persistencePath: '.atlas/rag',
        autoIndex: false,
        cacheEmbeddings: true
    };
    // Initialize RAG components
    const embeddingModel = new LocalEmbeddingModel({
        modelName: ragConfig.embeddingModel,
        cachePath: path.join(ragConfig.persistencePath || '.atlas/rag', 'embeddings-cache')
    });
    const vectorStore = new InMemoryVectorStore(embeddingModel);
    const documentProcessor = new MarkdownDocumentProcessor();
    // Create and initialize pipeline
    ragPipeline = new RAGPipeline(ragConfig, embeddingModel, vectorStore, documentProcessor);
    await ragPipeline.initialize();
    // Register request handler
    server.setRequestHandler(CallToolRequestSchema, async (request) => {
        const { name, arguments: args } = request.params;
        if (!args) {
            throw new Error('No arguments provided');
        }
        if (!ragPipeline) {
            throw new Error('RAG pipeline not initialized');
        }
        const projectPath = process.cwd();
        switch (name) {
            case 'rag_search': {
                const params = args;
                try {
                    const results = await ragPipeline.search(params);
                    if (results.length === 0) {
                        return {
                            content: [{
                                    type: 'text',
                                    text: 'No results found for your query.'
                                }]
                        };
                    }
                    const formattedResults = results.map((result, index) => {
                        const excerpt = result.chunk.content.substring(0, 200);
                        return `${index + 1}. **Score**: ${result.score.toFixed(3)}
   **Document**: ${result.document.path}
   **Content**: ${excerpt}${result.chunk.content.length > 200 ? '...' : ''}`;
                    }).join('\n\n');
                    return {
                        content: [{
                                type: 'text',
                                text: `=
 Found ${results.length} result${results.length > 1 ? 's' : ''}:\n\n${formattedResults}`
                            }]
                    };
                }
                catch (error) {
                    return {
                        content: [{
                                type: 'text',
                                text: `L Search failed: ${error}`
                            }]
                    };
                }
            }
            case 'rag_index_document': {
                const params = args;
                const documentPath = path.resolve(projectPath, params.path);
                try {
                    const document = await ragPipeline.indexDocument(documentPath);
                    return {
                        content: [{
                                type: 'text',
                                text: ` Successfully indexed document: ${params.path}\n\n` +
                                    `- Chunks created: ${document.chunks?.length || 0}\n` +
                                    `- Document ID: ${document.id}`
                            }]
                    };
                }
                catch (error) {
                    return {
                        content: [{
                                type: 'text',
                                text: `L Failed to index document: ${error}`
                            }]
                    };
                }
            }
            case 'rag_index_directory': {
                const params = args;
                const directoryPath = path.resolve(projectPath, params.path);
                try {
                    const result = await ragPipeline.indexDirectory(directoryPath);
                    return {
                        content: [{
                                type: 'text',
                                text: ` Successfully indexed ${result.indexed} documents from ${params.path}\n\n` +
                                    `- Failed: ${result.failed}\n` +
                                    (result.errors.length > 0 ? `- Errors:\n${result.errors.map(e => `  - ${e}`).join('\n')}` : '')
                            }]
                    };
                }
                catch (error) {
                    return {
                        content: [{
                                type: 'text',
                                text: `L Failed to index directory: ${error}`
                            }]
                    };
                }
            }
            case 'rag_index_collection': {
                const params = args;
                try {
                    const result = await ragPipeline.indexCollection(params.collection);
                    return {
                        content: [{
                                type: 'text',
                                text: ` Successfully indexed collection "${params.collection}"\n\n` +
                                    `- Documents indexed: ${result.indexed}\n` +
                                    `- Failed: ${result.failed}`
                            }]
                    };
                }
                catch (error) {
                    return {
                        content: [{
                                type: 'text',
                                text: `L Failed to index collection: ${error}`
                            }]
                    };
                }
            }
            case 'rag_get_stats': {
                try {
                    const stats = await ragPipeline.getStats();
                    return {
                        content: [{
                                type: 'text',
                                text: `=� RAG Statistics:\n\n` +
                                    `- Total Documents: ${stats.totalDocuments}\n` +
                                    `- Total Chunks: ${stats.totalChunks}\n` +
                                    `- Total Collections: ${stats.totalCollections}\n` +
                                    `- Index Size: ${(stats.indexSize / 1024 / 1024).toFixed(2)} MB\n` +
                                    `- Last Indexed: ${stats.lastIndexed}`
                            }]
                    };
                }
                catch (error) {
                    return {
                        content: [{
                                type: 'text',
                                text: `L Failed to get statistics: ${error}`
                            }]
                    };
                }
            }
            case 'rag_clear_index': {
                try {
                    await ragPipeline.clearIndex();
                    await ragPipeline.saveIndex();
                    return {
                        content: [{
                                type: 'text',
                                text: ` Successfully cleared the RAG index.`
                            }]
                    };
                }
                catch (error) {
                    return {
                        content: [{
                                type: 'text',
                                text: `L Failed to clear index: ${error}`
                            }]
                    };
                }
            }
            default:
                throw new Error(`Unknown tool: ${name}`);
        }
    });
    return {
        tools: [
            {
                name: 'rag_search',
                description: 'Search through indexed documents using semantic search',
                inputSchema: {
                    type: 'object',
                    properties: {
                        query: {
                            type: 'string',
                            description: 'Search query text'
                        },
                        limit: {
                            type: 'number',
                            description: 'Maximum number of results (default: 10)',
                            default: 10
                        },
                        threshold: {
                            type: 'number',
                            description: 'Minimum similarity threshold (0-1)',
                            minimum: 0,
                            maximum: 1
                        },
                        filters: {
                            type: 'object',
                            description: 'Metadata filters to apply',
                            additionalProperties: true
                        }
                    },
                    required: ['query']
                }
            },
            {
                name: 'rag_index_document',
                description: 'Index a single document for RAG',
                inputSchema: {
                    type: 'object',
                    properties: {
                        path: {
                            type: 'string',
                            description: 'Path to the document to index (relative to project root)'
                        }
                    },
                    required: ['path']
                }
            },
            {
                name: 'rag_index_directory',
                description: 'Index all markdown documents in a directory',
                inputSchema: {
                    type: 'object',
                    properties: {
                        path: {
                            type: 'string',
                            description: 'Path to the directory to index (relative to project root)'
                        }
                    },
                    required: ['path']
                }
            },
            {
                name: 'rag_index_collection',
                description: 'Index a predefined collection of documents',
                inputSchema: {
                    type: 'object',
                    properties: {
                        collection: {
                            type: 'string',
                            description: 'Name of the collection to index'
                        }
                    },
                    required: ['collection']
                }
            },
            {
                name: 'rag_get_stats',
                description: 'Get statistics about the RAG index',
                inputSchema: {
                    type: 'object',
                    properties: {},
                    required: []
                }
            },
            {
                name: 'rag_clear_index',
                description: 'Clear all indexed documents',
                inputSchema: {
                    type: 'object',
                    properties: {},
                    required: []
                }
            }
        ]
    };
}
//# sourceMappingURL=index.js.map