tools/code_analyzer.py

import ast
import astroid
from typing import Dict, List, Tuple, Optional
import re

class CodeAnalyzer:
    def __init__(self):
        self.algorithm_patterns = {
            'sort': ['bubble', 'quick', 'merge', 'insertion', 'selection'],
            'search': ['binary', 'linear', 'depth', 'breadth'],
            'graph': ['dijkstra', 'floyd', 'bellman', 'kruskal', 'prim'],
            'tree': ['binary', 'avl', 'red-black', 'b-tree'],
            'dynamic': ['knapsack', 'longest', 'shortest', 'matrix']
        }
        
        self.data_structure_patterns = {
            'linked_list': ['node', 'next', 'prev', 'head', 'tail'],
            'tree': ['node', 'left', 'right', 'root', 'leaf'],
            'graph': ['vertex', 'edge', 'adjacent', 'neighbor'],
            'stack': ['push', 'pop', 'peek', 'top'],
            'queue': ['enqueue', 'dequeue', 'front', 'rear'],
            'heap': ['heapify', 'sift', 'priority']
        }

    def analyze_code(self, code: str) -> Dict:
        """
        Analyze the given code and return information about its type and structure.
        """
        try:
            # Parse the code using ast
            tree = ast.parse(code)
            
            # Get basic code information
            info = {
                'type': 'unknown',
                'name': self._get_function_name(tree),
                'complexity': self._analyze_complexity(tree),
                'patterns': self._find_patterns(code),
                'structure': self._analyze_structure(tree)
            }
            
            # Determine the type of code
            info['type'] = self._determine_code_type(info['patterns'])
            
            return info
        except Exception as e:
            return {
                'type': 'error',
                'error': str(e)
            }

    def _get_function_name(self, tree: ast.AST) -> Optional[str]:
        """Extract the main function name from the code."""
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                return node.name
        return None

    def _analyze_complexity(self, tree: ast.AST) -> str:
        """Analyze the time complexity of the code."""
        # This is a simplified version - in practice, you'd want more sophisticated analysis
        loops = 0
        for node in ast.walk(tree):
            if isinstance(node, (ast.For, ast.While)):
                loops += 1
        
        if loops > 2:
            return "O(n³) or worse"
        elif loops == 2:
            return "O(n²)"
        elif loops == 1:
            return "O(n)"
        else:
            return "O(1)"

    def _find_patterns(self, code: str) -> Dict[str, List[str]]:
        """Find patterns in the code that indicate its type."""
        patterns = {
            'algorithms': [],
            'data_structures': []
        }
        
        # Check for algorithm patterns
        for category, keywords in self.algorithm_patterns.items():
            for keyword in keywords:
                if keyword.lower() in code.lower():
                    patterns['algorithms'].append(f"{category}:{keyword}")
        
        # Check for data structure patterns
        for structure, keywords in self.data_structure_patterns.items():
            for keyword in keywords:
                if keyword.lower() in code.lower():
                    patterns['data_structures'].append(f"{structure}:{keyword}")
        
        return patterns

    def _analyze_structure(self, tree: ast.AST) -> Dict:
        """Analyze the code structure and return relevant information."""
        structure = {
            'has_recursion': False,
            'has_classes': False,
            'has_functions': False,
            'has_loops': False
        }
        
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                structure['has_classes'] = True
            elif isinstance(node, ast.FunctionDef):
                structure['has_functions'] = True
                # Check for recursion: if the function calls itself
                for child in ast.walk(node):
                    if isinstance(child, ast.Call) and isinstance(child.func, ast.Name):
                        if child.func.id == node.name:
                            structure['has_recursion'] = True
            elif isinstance(node, (ast.For, ast.While)):
                structure['has_loops'] = True
        
        return structure

    def _determine_code_type(self, patterns: Dict) -> str:
        """Determine the type of code based on found patterns."""
        if patterns['algorithms']:
            return 'algorithm'
        elif patterns['data_structures']:
            return 'data_structure'
        else:
            return 'general'

    def get_explanation_prompt(self, code_info: Dict) -> str:
        """Generate a prompt for the LLM based on the code analysis."""
        prompt = f"This code appears to be a {code_info['type']} "
        
        if code_info['name']:
            prompt += f"named '{code_info['name']}' "
        
        prompt += f"with {code_info['complexity']} time complexity.\n\n"
        
        if code_info['patterns']['algorithms']:
            prompt += "It implements the following algorithms: " + ", ".join(code_info['patterns']['algorithms']) + ".\n"
        
        if code_info['patterns']['data_structures']:
            prompt += "It uses the following data structures: " + ", ".join(code_info['patterns']['data_structures']) + ".\n"
        
        prompt += "\nPlease explain how this code works step by step, focusing on:\n"
        prompt += "1. The main purpose and functionality\n"
        prompt += "2. The key components and how they work together\n"
        prompt += "3. The time and space complexity\n"
        prompt += "4. Any important edge cases or considerations\n"
        
        return prompt