src/services/code_graph_analyzer.py

"""
Code Graph Analyzer - Analyzes code relationships and dependencies
to provide better context for error analysis.
"""

import ast
import re
from typing import List, Dict, Set, Optional
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)

@dataclass
class FunctionInfo:
    """Information about a function"""
    name: str
    file_path: str
    line_number: int
    calls: List[str]  # Functions this function calls
    imports: List[str]  # Modules this function uses
    parameters: List[Dict[str, str]] = None  # List of {"name": name, "type": type}

@dataclass
class CodeGraphContext:
    """Enhanced context from code graph analysis"""
    related_functions: List[FunctionInfo]
    dependencies: List[str]
    execution_flow: List[str]

class CodeGraphAnalyzer:
    """Analyzes code structure and relationships"""
    
    def __init__(self):
        self.function_map: Dict[str, FunctionInfo] = {}
    
    def analyze_file(self, file_path: str, content: str) -> List[FunctionInfo]:
        """
        Parse a file and extract function information
        
        Args:
            file_path: Path to the file
            content: File content
            
        Returns:
            List of FunctionInfo objects
        """
        functions = []
        
        # Skip extremely large files (>500KB) or minified files
        file_size_kb = len(content) / 1024
        
        # Check if file is minified (very long lines, no formatting)
        lines = content.split('\n')
        if lines:
            avg_line_length = sum(len(line) for line in lines[:100]) / min(len(lines), 100)
            is_minified = avg_line_length > 200  # Minified files have very long lines
        else:
            is_minified = False
        
        # Skip if file is too large OR minified
        if file_size_kb > 500:
            logger.warning(f"Skipping graph analysis for very large file: {file_path} ({file_size_kb:.1f} KB)")
            return functions
        
        if is_minified:
            logger.warning(f"Skipping graph analysis for minified file: {file_path} (avg line length: {avg_line_length:.0f})")
            return functions
        
        try:
            # Parse the file based on extension
            if file_path.endswith(('.py',)):
                functions = self._analyze_python(file_path, content)
            elif file_path.endswith(('.js', '.jsx', '.ts', '.tsx')):
                functions = self._analyze_javascript(file_path, content)
            elif file_path.endswith('.cs'):
                functions = self._analyze_csharp(file_path, content)
            elif file_path.endswith('.java'):
                functions = self._analyze_java(file_path, content)
            elif file_path.endswith('.go'):
                functions = self._analyze_go(file_path, content)
            elif file_path.endswith('.php'):
                functions = self._analyze_php(file_path, content)
            elif file_path.endswith('.swift'):
                functions = self._analyze_swift(file_path, content)
            elif file_path.endswith('.kt'):
                functions = self._analyze_kotlin(file_path, content)
            elif file_path.endswith('.rs'):
                functions = self._analyze_rust(file_path, content)
            elif file_path.endswith(('.c', '.cpp', '.h', '.hpp')):
                functions = self._analyze_cpp(file_path, content)
            elif file_path.endswith('.rb'):
                functions = self._analyze_ruby(file_path, content)
            elif file_path.endswith('.scala'):
                functions = self._analyze_scala(file_path, content)
            elif file_path.endswith('.dart'):
                functions = self._analyze_dart(file_path, content)
            
        except Exception as e:
            logger.warning(f"Failed to analyze {file_path}: {e}")
        
        return functions
    
    def _analyze_python(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Python file using AST"""
        functions = []
        
        try:
            tree = ast.parse(content)
            
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    # Extract function calls and imports
                    calls = []
                    imports = []
                    parameters = []
                    
                    # Extract parameters and type hints
                    for arg in node.args.args:
                        arg_type = "Any"
                        if arg.annotation:
                            if isinstance(arg.annotation, ast.Name):
                                arg_type = arg.annotation.id
                            elif isinstance(arg.annotation, ast.Attribute):
                                arg_type = arg.annotation.attr
                        parameters.append({"name": arg.arg, "type": arg_type})

                    for child in ast.walk(node):
                        if isinstance(child, ast.Call):
                            if isinstance(child.func, ast.Name):
                                calls.append(child.func.id)
                            elif isinstance(child.func, ast.Attribute):
                                calls.append(child.func.attr)
                        
                        if isinstance(child, ast.Import):
                            for alias in child.names:
                                imports.append(alias.name)
                        elif isinstance(child, ast.ImportFrom):
                            if child.module:
                                imports.append(child.module)
                    
                    functions.append(FunctionInfo(
                        name=node.name,
                        file_path=file_path,
                        line_number=node.lineno,
                        calls=list(set(calls)),
                        imports=list(set(imports)),
                        parameters=parameters
                    ))
        
        except SyntaxError as e:
            logger.warning(f"Syntax error in {file_path}: {e}")
        
        return functions
    
    def _analyze_javascript(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze JavaScript/TypeScript file using regex (basic)"""
        functions = []
        
        # Simple regex patterns for function detection
        patterns = [
            r'function\s+(\w+)\s*\(',  # function name()
            r'const\s+(\w+)\s*=\s*\([^)]*\)\s*=>',  # const name = () =>
            r'(\w+)\s*:\s*function\s*\(',  # name: function()
        ]
        
        lines = content.split('\n')
        for i, line in enumerate(lines, 1):
            for pattern in patterns:
                match = re.search(pattern, line)
                if match:
                    func_name = match.group(1)
                    
                    # Extract function calls (simple heuristic)
                    calls = re.findall(r'(\w+)\s*\(', line)
                    calls = [c for c in calls if c != func_name]
                    
                    # Extract imports
                    imports = []
                    if 'import' in line:
                        import_match = re.search(r'from\s+[\'"]([^\'"]+)[\'"]', line)
                        if import_match:
                            imports.append(import_match.group(1))
                    
                    functions.append(FunctionInfo(
                        name=func_name,
                        file_path=file_path,
                        line_number=i,
                        calls=list(set(calls)),
                        imports=imports
                    ))
        
        return functions

    def _analyze_csharp(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze C# file using regex and brace counting"""
        # Improved pattern to capture method name and parameters
        method_pattern = r'(?:public|private|protected|internal|static|async|virtual|override|abstract|\s)*\s+(?:[\w<>\[\]?]+\s+)+(\w+)\s*\((.*?)\)'
        keywords = {'if', 'while', 'for', 'foreach', 'switch', 'catch', 'using', 'lock', 'fixed', 'checked', 'unchecked', 'typeof', 'sizeof', 'nameof', 'await', 'return', 'throw', 'new'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'using\s+([\w\.]+);')

    def _analyze_java(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Java file"""
        method_pattern = r'(?:public|private|protected|static|final|native|synchronized|abstract|transient|\s)*\s+[\w<>[\]]+\s+(\w+)\s*\((.*?)\)'
        keywords = {'if', 'while', 'for', 'switch', 'catch', 'synchronized', 'return', 'throw', 'new'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'import\s+([\w\.]+);')

    def _analyze_go(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Go file"""
        method_pattern = r'func\s+(?:\([\w\s\*]+\)\s+)?(\w+)\s*\((.*?)\)'
        keywords = {'if', 'for', 'switch', 'select', 'defer', 'go', 'return'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'import\s+"([\w\/]+)"')

    def _analyze_php(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze PHP file"""
        method_pattern = r'function\s+(\w+)\s*\((.*?)\)'
        keywords = {'if', 'while', 'for', 'foreach', 'switch', 'catch', 'return', 'echo', 'array'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'use\s+([\w\\]+);')

    def _analyze_swift(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Swift file"""
        method_pattern = r'func\s+(\w+)\s*\((.*?)\)'
        keywords = {'if', 'while', 'for', 'switch', 'catch', 'guard', 'defer', 'return', 'throw'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'import\s+(\w+)')

    def _analyze_kotlin(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Kotlin file"""
        method_pattern = r'fun\s+(\w+)\s*\((.*?)\)'
        keywords = {'if', 'while', 'for', 'when', 'catch', 'return', 'throw'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'import\s+([\w\.]+)')

    def _analyze_rust(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Rust file"""
        method_pattern = r'fn\s+(\w+)\s*\((.*?)\)'
        keywords = {'if', 'while', 'for', 'match', 'loop', 'return', 'break'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'use\s+([\w\:\{\}]+);')

    def _analyze_cpp(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze C/C++ file"""
        method_pattern = r'\w+\s+(\w+)\s*\((.*?)\)\s*\{'
        keywords = {'if', 'while', 'for', 'switch', 'catch', 'return', 'sizeof'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'#include\s*[<"]([\w\.]+)["\>]')

    def _analyze_ruby(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Ruby file (End-based, not brace-based)"""
        # Ruby requires a different parser because it uses 'end' instead of '}'
        functions = []
        method_pattern = r'def\s+(\w+)'
        keywords = {'if', 'unless', 'case', 'while', 'until', 'for', 'begin'} # These also use 'end'
        
        lines = content.split('\n')
        current_imports = []
        
        for i, line in enumerate(lines, 1):
            line_stripped = line.strip()
            if not line_stripped or line_stripped.startswith('#'):
                continue

            if 'require' in line_stripped:
                match = re.search(r'require\s+[\'"](.+)[\'"]', line_stripped)
                if match:
                    current_imports.append(match.group(1))
            
            match = re.search(method_pattern, line_stripped)
            if match:
                func_name = match.group(1)
                # Find calls on this line (pseudo-heuristic)
                calls = re.findall(r'(\w+)\(', line_stripped) # Ruby calls often don't have parens, but this is a safe start
                calls = [c for c in calls if c != func_name]
                
                functions.append(FunctionInfo(
                    name=func_name,
                    file_path=file_path,
                    line_number=i,
                    calls=list(set(calls)),
                    imports=list(current_imports)
                ))
        return functions

    def _analyze_scala(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Scala file"""
        method_pattern = r'def\s+(\w+)\s*[\(\[]'
        keywords = {'if', 'while', 'for', 'match', 'return', 'throw', 'new', 'try', 'catch'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'import\s+([\w\.]+)')

    def _analyze_dart(self, file_path: str, content: str) -> List[FunctionInfo]:
        """Analyze Dart file"""
        # ReturnType Name(Args)
        method_pattern = r'[\w<>]+\s+(\w+)\s*\(' 
        keywords = {'if', 'while', 'for', 'switch', 'return', 'throw', 'new', 'try', 'catch', 'await'}
        return self._generic_brace_parser(file_path, content, method_pattern, keywords, import_pattern=r'import\s+[\'"](.+)[\'"];')

    def _generic_brace_parser(self, file_path: str, content: str, method_pattern: str, keywords: Set[str], import_pattern: str = None) -> List[FunctionInfo]:
        """Generic parser for brace-based languages (C#, Java, Go, etc.)"""
        functions = []
        lines = content.split('\n')
        current_imports = []
        
        brace_balance = 0
        active_function = None 
        function_start_balance = 0
        
        for i, line in enumerate(lines, 1):
            line_stripped = line.strip()
            if not line_stripped:
                continue
            
            # Imports
            if import_pattern and re.search(import_pattern, line_stripped):
                match = re.search(import_pattern, line_stripped)
                if match:
                    current_imports.append(match.group(1))

            # Method Start
            if not active_function:
                 if any(line_stripped.startswith(kw + ' ') or line_stripped.startswith(kw + '(') for kw in keywords):
                    pass
                 else:
                    match = re.search(method_pattern, line_stripped)
                    if match:
                        func_name = match.group(1)
                        param_str = match.group(2) if match.lastindex >= 2 else ""
                        
                        parameters = []
                        if param_str:
                            # Basic parameter splitting for C-style languages
                            params = param_str.split(',')
                            for p in params:
                                p = p.strip()
                                if not p: continue
                                # Split on space: "Type name" or "Type? name"
                                p_parts = p.split()
                                if len(p_parts) >= 2:
                                    # Last part is likely name, everything before is type
                                    name = p_parts[-1]
                                    p_type = " ".join(p_parts[:-1])
                                    parameters.append({"name": name, "type": p_type})
                                else:
                                    parameters.append({"name": p, "type": "Any"})

                        if func_name not in keywords:
                            active_function = FunctionInfo(
                                name=func_name,
                                file_path=file_path,
                                line_number=i,
                                calls=[],
                                imports=list(current_imports),
                                parameters=parameters
                            )
                            function_start_balance = brace_balance

            # Calls
            if active_function:
                calls = re.findall(r'(\w+)\s*\(', line_stripped)
                filtered_calls = [c for c in calls if c != active_function.name and c not in keywords]
                active_function.calls.extend(filtered_calls)
            
            # Balance Support (simplified for string safety)
            # Remove string contents before counting braces to avoid "{}" inside strings
            line_no_strings = re.sub(r'".*?"|\'.*?\'', '', line)
            open_braces = line_no_strings.count('{')
            close_braces = line_no_strings.count('}')
            brace_balance += (open_braces - close_braces)
            
            # Function End
            if active_function:
                if brace_balance <= function_start_balance and (open_braces > 0 or close_braces > 0 or ';' in line):
                    # Handle one-liners or end of block
                    # Only close if we actually saw a brace change related to closure or started with implicit closure
                     if brace_balance <= function_start_balance and (open_braces > 0 or close_braces > 0):
                         functions.append(active_function)
                         active_function = None
                     elif ';' in line and brace_balance == function_start_balance and open_braces == 0:
                         # Abstract/ Interface method ending
                         functions.append(active_function)
                         active_function = None

        for f in functions:
            f.calls = list(set(f.calls))
        return functions
    
    def build_graph(self, files: List[Dict[str, str]]) -> None:
        """
        Build function graph from multiple files
        
        Args:
            files: List of dicts with 'path' and 'content' keys
        """
        self.function_map = {}
        
        for file in files:
            functions = self.analyze_file(file['path'], file['content'])
            for func in functions:
                key = f"{file['path']}:{func.name}"
                self.function_map[key] = func
                
    def save_graph(self, collection_name: str, output_dir: str = "./repo_metadata") -> None:
        """Save the built graph to a JSON file"""
        import json
        import os
        from dataclasses import asdict
        
        os.makedirs(output_dir, exist_ok=True)
        file_path = os.path.join(output_dir, f"{collection_name}_graph.json")
        
        try:
            # Convert FunctionInfo objects to dicts
            data = {k: asdict(v) for k, v in self.function_map.items()}
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2)
            logger.info(f"Saved code graph to {file_path}")
        except Exception as e:
            logger.error(f"Failed to save code graph: {e}")

    def load_graph(self, collection_name: str, input_dir: str = "./repo_metadata") -> bool:
        """Load a graph from a JSON file and merge it into current graph"""
        import json
        import os
        
        file_path = os.path.join(input_dir, f"{collection_name}_graph.json")
        if not os.path.exists(file_path):
            logger.warning(f"No existing graph found at {file_path}")
            return False
            
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            if not hasattr(self, 'function_map'):
                self.function_map = {}
                
            loaded_count = 0
            for k, v in data.items():
                self.function_map[k] = FunctionInfo(**v)
                loaded_count += 1
                
            logger.info(f"Loaded/Merged code graph from {file_path} ({loaded_count} functions)")
            return True
        except Exception as e:
            logger.error(f"Failed to load code graph: {e}")
            return False

    def find_related_functions(self, error_location: str, max_depth: int = 2) -> List[FunctionInfo]:
        """
        Find functions related to the error location
        """
        related = []
        visited = set()
        
        try:
            if ':' not in error_location:
                return []
                
            file_path, line_str = error_location.rsplit(':', 1)
            try:
                line_num = int(line_str)
            except ValueError:
                return []
            
            # Find function at this location
            start_func = None
            for key, func in self.function_map.items():
                if func.file_path.endswith(file_path.split('/')[-1]) and func.file_path in file_path:
                    if func.line_number <= line_num:
                        if start_func is None or func.line_number > start_func.line_number:
                            start_func = func
            
            if start_func:
                logger.debug(f"Graph traversal starting at {start_func.name}")
                self._traverse_graph(start_func, related, visited, max_depth)
                
                # Pillar 3: Data-Flow Enrichment
                # If we have parameters with types, let's find their definitions
                additional_definitions = []
                for func in related:
                    if func.parameters:
                        for param in func.parameters:
                            p_type = param.get('type')
                            if p_type and p_type not in ('Any', 'int', 'str', 'bool', 'Dict', 'List', 'None'):
                                # Try to find where this class/type is defined
                                type_def_func = self.find_type_definition(p_type)
                                if type_def_func and type_def_func not in related:
                                    additional_definitions.append(type_def_func)
                
                related.extend(additional_definitions)
        
        except Exception as e:
            logger.warning(f"Error finding related functions: {e}")
        
        return related

    def find_type_definition(self, type_name: str) -> Optional[FunctionInfo]:
        """
        Try to find the definition of a type (class/struct) in the function map.
        In this simplified graph, we check for functions with the same name (constructors)
        or classes (not yet explicitly tracked, so we look for common patterns).
        """
        # 1. Look for a function with this name (likely a constructor or the class itself in some parsers)
        for key, func in self.function_map.items():
            if func.name == type_name:
                return func
                
        return None
    
    def resolve_import_path(self, import_name: str, current_file_path: str) -> List[str]:
        """
        Resolve an import statement to potential file paths in the graph.
        e.g., 'src.services.utils' -> ['src/services/utils.py']
        """
        potential_paths = []
        
        # 1. Direct mapping (Python style dots to slashes)
        base_path = import_name.replace('.', '/')
        
        # Try common extensions
        extensions = ['.py', '.js', '.ts', '.cs', '.java', '.go']
        for ext in extensions:
             potential_paths.append(base_path + ext)
             potential_paths.append(f"{base_path}/__init__{ext}") # Python package
             potential_paths.append(f"{base_path}/index{ext}")    # JS package
        
        # 2. Relative imports (starting with .)
        if import_name.startswith('.'):
            # Resolve relative to current file
            current_dir = '/'.join(current_file_path.split('/')[:-1])
            relative_base = base_path.lstrip('./')
            combined = f"{current_dir}/{relative_base}"
            for ext in extensions:
                potential_paths.append(combined + ext)

        # 3. Filter by what actually exists in our function map
        resolved_files = []
        known_files = set()
        for key in self.function_map.keys():
            known_files.add(key.split(':')[0])
            
        for path in potential_paths:
            # We check if any known file ENDS with this path (partial match for flexibility)
            for known in known_files:
                if known.endswith(path):
                    resolved_files.append(known)
        
        return list(set(resolved_files))

    def _resolve_call_to_function(self, call_name: str, caller_func: FunctionInfo) -> List[FunctionInfo]:
        """
        Find the FunctionInfo that matches a function call, checking imports.
        """
        candidates = []
        
        # 1. Check local file first (same class/module)
        local_key = f"{caller_func.file_path}:{call_name}"
        if local_key in self.function_map:
            candidates.append(self.function_map[local_key])
        
        # 2. Check if call is "module.func"
        if '.' in call_name:
            module_part, func_part = call_name.rsplit('.', 1)
            
            # Check implicit imports (fields/properties) or explicit imports
            for imp in caller_func.imports:
                # If we imported the module 'src.services.utils'
                # And call is 'utils.helper' -> module_part 'utils' matches end of import
                if imp.endswith(module_part) or imp == module_part:
                    # Resolve this import to a file
                    resolved_files = self.resolve_import_path(imp, caller_func.file_path)
                    for file_path in resolved_files:
                        target_key = f"{file_path}:{func_part}"
                        if target_key in self.function_map:
                            candidates.append(self.function_map[target_key])
                            
        # 3. Check explicit "from X import Y"
        # If call_name is just 'Y', and we have 'from X import Y'
        # This is harder to track without full AST symbol tables, 
        # but we can check if 'call_name' appears in any resolved file's functions
        else:
            for imp in caller_func.imports:
                resolved_files = self.resolve_import_path(imp, caller_func.file_path)
                for file_path in resolved_files:
                    target_key = f"{file_path}:{call_name}"
                    if target_key in self.function_map:
                        candidates.append(self.function_map[target_key])
                        
        return candidates

    def _traverse_graph(self, func: FunctionInfo, related: List[FunctionInfo], 
                       visited: Set[str], depth: int) -> None:
        """Recursively traverse function call graph with Cross-File resolution"""
        if depth <= 0:
            return
        
        key = f"{func.file_path}:{func.name}"
        if key in visited:
            return
        
        visited.add(key)
        related.append(func)
        
        # Find functions this one calls
        for call in func.calls:
            # Resolve the call to actual function definitions
            targets = self._resolve_call_to_function(call, func)
            
            # Use simple name matching fallback if resolution failed and it's a simple name
            if not targets and '.' not in call:
                for map_key, called_func in self.function_map.items():
                    if called_func.name == call and called_func.file_path == func.file_path: # Prefer local
                         targets.append(called_func)
            
            for target_func in targets:
                self._traverse_graph(target_func, related, visited, depth - 1)
    
    def get_execution_flow(self, error_location: str) -> List[str]:
        """
        Get likely execution flow leading to error
        
        Args:
            error_location: File path and line number
            
        Returns:
            List of function names in execution order
        """
        related = self.find_related_functions(error_location)
        return [f"{func.file_path}:{func.name}()" for func in related]
    
    def get_dependencies(self, error_location: str) -> List[str]:
        """
        Get all dependencies related to error location
        
        Args:
            error_location: File path and line number
            
        Returns:
            List of import/dependency names
        """
        related = self.find_related_functions(error_location)
        all_imports = set()
        
        for func in related:
            all_imports.update(func.imports)
        
        return list(all_imports)