import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { ConfigManager } from '../../config/config-manager.js';
import { CodeEmbeddings } from './embeddings.js';
import { HybridCodeMemory } from './graph-memory.js';
import { promises as fs } from 'fs';
import path from 'path';

let codeEmbeddings: CodeEmbeddings;
let codeMemory: HybridCodeMemory;
let configManager: ConfigManager;

export async function setupLocalAITools(server: Server, config?: ConfigManager) {
  configManager = config || new ConfigManager();
  codeEmbeddings = new CodeEmbeddings();
  codeMemory = new HybridCodeMemory();

  // Add local AI tool handlers
  server.setRequestHandler(CallToolRequestSchema, async (request) => {
    const { name, arguments: args } = request.params;

    if (!args) {
      throw new Error('No arguments provided');
    }

    switch (name) {
      case 'index_codebase': {
        const { directory = '.', extensions = ['.ts', '.js', '.tsx', '.jsx', '.py'] } = args as {
          directory?: string;
          extensions?: string[];
        };
        
        const dir = path.isAbsolute(directory) 
          ? directory 
          : path.join(process.cwd(), directory);
        
        try {
          const files = await findCodeFiles(dir, extensions);
          let indexed = 0;
          
          for (const file of files) {
            try {
              const content = await fs.readFile(file, 'utf-8');
              const language = detectLanguage(file);
              
              // Add to embeddings for semantic search
              codeEmbeddings.addCodeFile(file, content, language);
              
              // Add to graph memory
              const relativePath = path.relative(dir, file);
              codeMemory.recordCodeChange({
                file: relativePath,
                type: 'create',
                functions: extractFunctions(content, language),
                imports: extractImports(content, language),
              });
              
              indexed++;
            } catch (error) {
              // Skip files that can't be read
            }
          }
          
          return {
            content: [
              {
                type: 'text',
                text: `🔍 Indexed ${indexed} files for semantic search and analysis\n\n` +
                      `📊 Total embeddings: ${codeEmbeddings.size()}\n` +
                      `🧠 Memory initialized with code relationships`,
              },
            ],
            structured: {
              indexed,
              total: files.length,
              embeddingCount: codeEmbeddings.size(),
            },
          };
        } catch (error) {
          return {
            content: [
              {
                type: 'text',
                text: `❌ Error indexing codebase: ${(error as Error).message}`,
              },
            ],
          };
        }
      }
      
      case 'semantic_search': {
        const { query, topK = 5 } = args as { query: string; topK?: number };
        
        const results = codeEmbeddings.searchCode(query, topK);
        
        if (results.length === 0) {
          return {
            content: [
              {
                type: 'text',
                text: `No results found for: "${query}"\n\nTry indexing the codebase first with 'index_codebase'`,
              },
            ],
          };
        }
        
        const formattedResults = results.map((result, i) => {
          const preview = result.text.split('\n').slice(0, 5).join('\n');
          return `${i + 1}. ${result.metadata?.filePath || 'Unknown'} (score: ${result.score.toFixed(3)})\n${preview}...`;
        }).join('\n\n');
        
        return {
          content: [
            {
              type: 'text',
              text: `🔍 Semantic search results for: "${query}"\n\n${formattedResults}`,
            },
          ],
          structured: {
            query,
            results: results.map(r => ({
              file: r.metadata?.filePath,
              score: r.score,
              features: r.metadata?.features,
            })),
          },
        };
      }
      
      case 'analyze_dependencies': {
        const { file } = args as { file: string };
        
        const insights = codeMemory.getInsights();
        const fileQuery = {
          type: 'graph' as const,
          query: `file:${file}`,
          depth: 2,
        };
        
        const dependencies = codeMemory.query(fileQuery) as any[];
        
        return {
          content: [
            {
              type: 'text',
              text: `📊 Dependency Analysis for ${file}\n\n` +
                    `🔗 Direct dependencies: ${dependencies.length}\n` +
                    `🔄 Circular dependencies found: ${insights.dependencies.length}\n` +
                    `🔥 Code hotspots:\n${insights.hotspots.slice(0, 5).join('\n')}\n\n` +
                    `📈 Recent patterns:\n${Object.entries(insights.recentPatterns)
                      .map(([k, v]) => `  ${k}: ${v} occurrences`)
                      .join('\n')}`,
            },
          ],
          structured: {
            file,
            dependencies,
            insights,
          },
        };
      }
      
      case 'find_similar_code': {
        const { filePath } = args as { filePath: string };
        
        try {
          const content = await fs.readFile(filePath, 'utf-8');
          const results = codeEmbeddings.searchCode(content, 5);
          
          // Filter out the same file
          const similar = results.filter(r => r.metadata?.filePath !== filePath);
          
          if (similar.length === 0) {
            return {
              content: [
                {
                  type: 'text',
                  text: `No similar code found. Try indexing more files first.`,
                },
              ],
            };
          }
          
          const formatted = similar.map((result, i) => 
            `${i + 1}. ${result.metadata?.filePath} (similarity: ${(result.score * 100).toFixed(1)}%)`
          ).join('\n');
          
          return {
            content: [
              {
                type: 'text',
                text: `🔍 Files similar to ${path.basename(filePath)}:\n\n${formatted}\n\n` +
                    `💡 These files might have:\n` +
                    `• Similar functionality\n` +
                    `• Reusable patterns\n` +
                    `• Potential for consolidation`,
              },
            ],
            structured: {
              sourceFile: filePath,
              similarFiles: similar.map(r => ({
                file: r.metadata?.filePath,
                similarity: r.score,
              })),
            },
          };
        } catch (error) {
          return {
            content: [
              {
                type: 'text',
                text: `❌ Error analyzing file: ${(error as Error).message}`,
              },
            ],
          };
        }
      }
      
      case 'track_code_change': {
        const { file, changeType, content } = args as {
          file: string;
          changeType: 'create' | 'modify' | 'delete';
          content?: string;
        };
        
        if (content) {
          const language = detectLanguage(file);
          
          // Update embeddings
          if (changeType !== 'delete') {
            codeEmbeddings.addCodeFile(file, content, language);
          }
          
          // Record in memory
          codeMemory.recordCodeChange({
            file,
            type: changeType,
            functions: extractFunctions(content, language),
            imports: extractImports(content, language),
          });
        }
        
        return {
          content: [
            {
              type: 'text',
              text: `✅ Tracked ${changeType} for ${file}\n` +
                    `🧠 Memory updated with new relationships`,
            },
          ],
        };
      }
      
      default:
        // Not a local AI tool
        return null as any;
    }
  });

  return {
    tools: [
      {
        name: 'index_codebase',
        description: 'Index codebase for semantic search and dependency analysis',
        inputSchema: {
          type: 'object',
          properties: {
            directory: { type: 'string', default: '.', description: 'Directory to index' },
            extensions: { 
              type: 'array', 
              items: { type: 'string' },
              default: ['.ts', '.js', '.tsx', '.jsx', '.py'],
              description: 'File extensions to include'
            },
          },
          required: [],
        },
      },
      {
        name: 'semantic_search',
        description: 'Search codebase using natural language queries',
        inputSchema: {
          type: 'object',
          properties: {
            query: { type: 'string', description: 'Natural language search query' },
            topK: { type: 'number', default: 5, description: 'Number of results to return' },
          },
          required: ['query'],
        },
      },
      {
        name: 'find_similar_code',
        description: 'Find files with similar code patterns',
        inputSchema: {
          type: 'object',
          properties: {
            filePath: { type: 'string', description: 'Path to file to find similar code for' },
          },
          required: ['filePath'],
        },
      },
      {
        name: 'analyze_dependencies',
        description: 'Analyze code dependencies and relationships',
        inputSchema: {
          type: 'object',
          properties: {
            file: { type: 'string', description: 'File to analyze dependencies for' },
          },
          required: ['file'],
        },
      },
      {
        name: 'track_code_change',
        description: 'Track code changes in memory for pattern analysis',
        inputSchema: {
          type: 'object',
          properties: {
            file: { type: 'string', description: 'File that changed' },
            changeType: { 
              type: 'string', 
              enum: ['create', 'modify', 'delete'],
              description: 'Type of change'
            },
            content: { type: 'string', description: 'New content (if not delete)' },
          },
          required: ['file', 'changeType'],
        },
      },
    ],
    resources: [],
  };
}

async function findCodeFiles(dir: string, extensions: string[]): Promise<string[]> {
  const files: string[] = [];
  
  async function walk(currentDir: string) {
    const entries = await fs.readdir(currentDir, { withFileTypes: true });
    
    for (const entry of entries) {
      const fullPath = path.join(currentDir, entry.name);
      
      if (entry.isDirectory() && !entry.name.startsWith('.') && entry.name !== 'node_modules') {
        await walk(fullPath);
      } else if (entry.isFile() && extensions.some(ext => entry.name.endsWith(ext))) {
        files.push(fullPath);
      }
    }
  }
  
  await walk(dir);
  return files;
}

function detectLanguage(filePath: string): string {
  const ext = path.extname(filePath).toLowerCase();
  const languageMap: Record<string, string> = {
    '.ts': 'typescript',
    '.tsx': 'typescript',
    '.js': 'javascript',
    '.jsx': 'javascript',
    '.py': 'python',
    '.rs': 'rust',
    '.go': 'go',
  };
  
  return languageMap[ext] || 'unknown';
}

function extractFunctions(content: string, language: string): string[] {
  const functions: string[] = [];
  
  if (language === 'typescript' || language === 'javascript') {
    const functionRegex = /(?:function|const|let|var)\s+(\w+)\s*(?:=\s*)?(?:\([^)]*\)|async)/g;
    let match;
    while ((match = functionRegex.exec(content)) !== null) {
      functions.push(match[1]);
    }
  }
  
  return functions;
}

function extractImports(content: string, language: string): string[] {
  const imports: string[] = [];
  
  if (language === 'typescript' || language === 'javascript') {
    const importRegex = /import\s+.*\s+from\s+['"]([^'"]+)['"]/g;
    let match;
    while ((match = importRegex.exec(content)) !== null) {
      imports.push(match[1]);
    }
  }
  
  return imports;
}