// code-context-server.js - A resource that understands your code

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import * as ts from 'typescript';
import * as babel from '@babel/parser';
import path from 'path';
import fs from 'fs/promises';

class CodeContextProvider {
  constructor(projectRoot) {
    this.projectRoot = projectRoot;
    this.cache = new Map();
  }
  
  // Provide rich context about a code file
  async getFileContext(filePath) {
    const content = await fs.readFile(filePath, 'utf-8');
    const ext = path.extname(filePath);
    
    let context = {
      uri: `file://${filePath}`,
      content: content,
      metadata: {
        lines: content.split('\n').length,
        size: content.length,
        lastModified: (await fs.stat(filePath)).mtime
      }
    };
    
    // Language-specific analysis
    if (['.js', '.jsx', '.ts', '.tsx'].includes(ext)) {
      context.analysis = await this.analyzeJavaScript(content, filePath);
    } else if (['.py'].includes(ext)) {
      context.analysis = await this.analyzePython(content);
    }
    
    // Find related files
    context.related = await this.findRelatedFiles(filePath);
    
    // Extract TODOs and FIXMEs
    context.annotations = this.extractAnnotations(content);
    
    return context;
  }
  
  async analyzeJavaScript(content, filePath) {
    const isTypeScript = filePath.endsWith('.ts') || filePath.endsWith('.tsx');
    
    try {
      const ast = babel.parse(content, {
        sourceType: 'module',
        plugins: [
          'jsx',
          isTypeScript && 'typescript',
          'classProperties',
          'asyncGenerators',
          'objectRestSpread'
        ].filter(Boolean)
      });
      
      const analysis = {
        imports: [],
        exports: [],
        functions: [],
        classes: [],
        complexity: 0
      };
      
      // Traverse AST to extract structure
      babel.traverse(ast, {
        ImportDeclaration(path) {
          analysis.imports.push({
            source: path.node.source.value,
            specifiers: path.node.specifiers.map(spec => spec.local.name)
          });
        },
        
        ExportNamedDeclaration(path) {
          if (path.node.declaration) {
            analysis.exports.push({
              type: 'named',
              name: path.node.declaration.id?.name || 'anonymous'
            });
          }
        },
        
        FunctionDeclaration(path) {
          analysis.functions.push({
            name: path.node.id?.name || 'anonymous',
            params: path.node.params.length,
            async: path.node.async,
            lines: path.node.loc.end.line - path.node.loc.start.line
          });
          analysis.complexity += this.calculateComplexity(path);
        },
        
        ClassDeclaration(path) {
          analysis.classes.push({
            name: path.node.id?.name || 'anonymous',
            methods: path.node.body.body
              .filter(member => member.type === 'ClassMethod')
              .map(method => method.key.name)
          });
        }
      });
      
      return analysis;
    } catch (error) {
      console.error('Failed to parse JavaScript:', error);
      return { parseError: error.message };
    }
  }
  
  async findRelatedFiles(filePath) {
    const related = [];
    const fileName = path.basename(filePath, path.extname(filePath));
    const dirPath = path.dirname(filePath);
    
    // Look for test files
    const testPatterns = [
      `${fileName}.test.*`,
      `${fileName}.spec.*`,
      `__tests__/${fileName}.*`
    ];
    
    // Look for related files (same name, different extension)
    const files = await fs.readdir(dirPath);
    
    for (const file of files) {
      const baseName = path.basename(file, path.extname(file));
      
      // Test files
      if (file.includes('.test.') || file.includes('.spec.')) {
        if (baseName.includes(fileName)) {
          related.push({
            type: 'test',
            uri: `file://${path.join(dirPath, file)}`
          });
        }
      }
      
      // Style files
      if (baseName === fileName && ['.css', '.scss', '.module.css'].includes(path.extname(file))) {
        related.push({
          type: 'style',
          uri: `file://${path.join(dirPath, file)}`
        });
      }
      
      // Component files
      if (baseName === fileName && file !== path.basename(filePath)) {
        related.push({
          type: 'related',
          uri: `file://${path.join(dirPath, file)}`
        });
      }
    }
    
    return related;
  }
  
  extractAnnotations(content) {
    const annotations = [];
    const lines = content.split('\n');
    
    const patterns = [
      { pattern: /\/\/\s*(TODO|FIXME|HACK|XXX|NOTE):?\s*(.+)/, type: 'comment' },
      { pattern: /#\s*(TODO|FIXME|HACK|XXX|NOTE):?\s*(.+)/, type: 'comment' },
      { pattern: /@(todo|deprecated|throws|param|returns)/, type: 'jsdoc' }
    ];
    
    lines.forEach((line, index) => {
      for (const { pattern, type } of patterns) {
        const match = line.match(pattern);
        if (match) {
          annotations.push({
            line: index + 1,
            type: type,
            tag: match[1],
            text: match[2] || line.trim(),
            severity: match[1] === 'FIXME' ? 'high' : 'normal'
          });
        }
      }
    });
    
    return annotations;
  }
  
  calculateComplexity(path) {
    let complexity = 1;
    
    path.traverse({
      IfStatement() { complexity++; },
      ForStatement() { complexity++; },
      WhileStatement() { complexity++; },
      DoWhileStatement() { complexity++; },
      SwitchCase() { complexity++; },
      CatchClause() { complexity++; },
      ConditionalExpression() { complexity++; },
      LogicalExpression({ node }) {
        if (node.operator === '&&' || node.operator === '||') {
          complexity++;
        }
      }
    });
    
    return complexity;
  }
}

// MCP Server implementation
const server = new Server(
  {
    name: 'code-context',
    version: '1.0.0',
  },
  {
    capabilities: {
      resources: {}
    }
  }
);

const projectRoot = process.argv[2] || process.cwd();
const provider = new CodeContextProvider(projectRoot);

// List available resources
server.setRequestHandler('resources/list', async () => {
  // Scan for code files
  const codeFiles = await scanForCodeFiles(projectRoot);
  
  return {
    resources: [
      {
        uri: `code:///${projectRoot}`,
        name: 'Project Overview',
        mimeType: 'application/json',
        description: 'Overview of the entire project structure'
      },
      ...codeFiles.slice(0, 100).map(file => ({
        uri: `code://file${file}`,
        name: path.relative(projectRoot, file),
        mimeType: 'application/json',
        description: 'Rich context about this code file'
      })),
      {
        uri: `code:///search`,
        name: 'Code Search',
        mimeType: 'application/json',
        description: 'Search across all code files'
      }
    ]
  };
});

// Read specific resources
server.setRequestHandler('resources/read', async (request) => {
  const { uri } = request.params;
  const parsed = new URL(uri);
  
  if (parsed.pathname === '/' && parsed.hostname === '') {
    // Project overview
    const overview = await generateProjectOverview(projectRoot);
    
    return {
      contents: [{
        uri: uri,
        mimeType: 'application/json',
        text: JSON.stringify(overview, null, 2)
      }]
    };
  }
  
  if (parsed.pathname.startsWith('/file')) {
    // Specific file context
    const filePath = parsed.pathname.substring(5); // Remove '/file' prefix
    const context = await provider.getFileContext(filePath);
    
    return {
      contents: [{
        uri: uri,
        mimeType: 'application/json', 
        text: JSON.stringify(context, null, 2)
      }]
    };
  }
  
  if (parsed.pathname === '/search') {
    // Search functionality
    const query = parsed.searchParams.get('q');
    const type = parsed.searchParams.get('type') || 'all';
    
    const results = await searchCode(projectRoot, query, type);
    
    return {
      contents: [{
        uri: uri,
        mimeType: 'application/json',
        text: JSON.stringify(results, null, 2)
      }]
    };
  }
  
  throw new Error(`Unknown resource: ${uri}`);
});

// Helper functions
async function scanForCodeFiles(dir, files = []) {
  const entries = await fs.readdir(dir, { withFileTypes: true });
  
  for (const entry of entries) {
    const fullPath = path.join(dir, entry.name);
    
    if (entry.isDirectory()) {
      // Skip common directories
      if (!['node_modules', '.git', 'dist', 'build'].includes(entry.name)) {
        await scanForCodeFiles(fullPath, files);
      }
    } else if (entry.isFile()) {
      // Include code files
      if (/\.(js|jsx|ts|tsx|py|java|cpp|c|h|go|rs|rb|php)$/.test(entry.name)) {
        files.push(fullPath);
      }
    }
  }
  
  return files;
}

async function generateProjectOverview(projectRoot) {
  const files = await scanForCodeFiles(projectRoot);
  
  const overview = {
    root: projectRoot,
    statistics: {
      totalFiles: files.length,
      byLanguage: {},
      totalLines: 0
    },
    structure: {},
    dependencies: await analyzeDependencies(projectRoot),
    recentActivity: await getRecentActivity(projectRoot)
  };
  
  // Analyze files
  for (const file of files) {
    const ext = path.extname(file);
    const content = await fs.readFile(file, 'utf-8');
    const lines = content.split('\n').length;
    
    overview.statistics.totalLines += lines;
    overview.statistics.byLanguage[ext] = (overview.statistics.byLanguage[ext] || 0) + 1;
  }
  
  return overview;
}

const transport = new StdioServerTransport();
await server.connect(transport);
console.error(`Code Context Server running for: ${projectRoot}`);