import ast
import json
from typing import Any, Dict, List, Set
from pathlib import Path
from mcp.server.fastmcp import FastMCP

# Create FastMCP server
mcp = FastMCP("code_analyzer")

@mcp.tool()
def analyze_structure(file_path: str) -> str:
    """
    Analyze Python code structure (functions and classes).

    Args:
        file_path: Path to the Python file to analyze

    Returns:
        JSON string containing structure information
    """
    path = Path(file_path)
    if not path.exists():
        return json.dumps({"error": f"File not found: {file_path}"})

    try:
        with open(path, 'r') as f:
            content = f.read()
        tree = ast.parse(content)

        functions = []
        classes = []

        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                functions.append({
                    'name': node.name,
                    'arguments': [arg.arg for arg in node.args.args],
                    'line_number': node.lineno
                })
            elif isinstance(node, ast.ClassDef):
                methods = [n.name for n in ast.walk(node) if isinstance(n, ast.FunctionDef)]
                classes.append({
                    'name': node.name,
                    'methods': methods,
                    'line_number': node.lineno
                })

        return json.dumps({
            "type": "structure",
            "functions": functions,
            "classes": classes
        }, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})

@mcp.tool()
def analyze_complexity(file_path: str) -> str:
    """
    Analyze Python code complexity (cyclomatic complexity).

    Args:
        file_path: Path to the Python file to analyze

    Returns:
        JSON string containing complexity information
    """
    path = Path(file_path)
    if not path.exists():
        return json.dumps({"error": f"File not found: {file_path}"})

    try:
        with open(path, 'r') as f:
            content = f.read()
        tree = ast.parse(content)

        complexity_info = {}

        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                decisions = [n for n in ast.walk(node) if isinstance(n, (
                    ast.If, ast.For, ast.While, ast.And, ast.Or, ast.ExceptHandler
                ))]

                returns = len([n for n in ast.walk(node) if isinstance(n, ast.Return)])
                breaks = len([n for n in ast.walk(node) if isinstance(n, ast.Break)])

                cyclomatic_complexity = len(decisions) + 1
                variables = _get_variable_usage(node)

                complexity_info[node.name] = {
                    'cyclomatic_complexity': cyclomatic_complexity,
                    'complexity_level': _get_complexity_level(cyclomatic_complexity),
                    'decision_points': len(decisions),
                    'returns': returns,
                    'breaks': breaks,
                    'line_count': node.end_lineno - node.lineno if hasattr(node, 'end_lineno') else 1,
                    'variable_uses': len(variables),
                    'variables': list(variables)
                }

        return json.dumps({
            "type": "complexity",
            "functions": complexity_info,
            "summary": {
                "avg_complexity": sum(f['cyclomatic_complexity'] for f in complexity_info.values()) / len(complexity_info) if complexity_info else 0,
                "most_complex_function": max(complexity_info.items(), key=lambda x: x[1]['cyclomatic_complexity'])[0] if complexity_info else None
            }
        }, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})

@mcp.tool()
def analyze_dependencies(file_path: str) -> str:
    """
    Analyze Python code dependencies (imports).

    Args:
        file_path: Path to the Python file to analyze

    Returns:
        JSON string containing dependency information
    """
    path = Path(file_path)
    if not path.exists():
        return json.dumps({"error": f"File not found: {file_path}"})

    try:
        with open(path, 'r') as f:
            content = f.read()
        tree = ast.parse(content)

        imports = []
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                imports.extend(n.name for n in node.names)
            elif isinstance(node, ast.ImportFrom):
                imports.append(f"{node.module}.{node.names[0].name}")

        return json.dumps({
            "type": "dependencies",
            "imports": imports
        }, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})

def _get_complexity_level(complexity: int) -> str:
    if complexity <= 5:
        return "low"
    elif complexity <= 10:
        return "moderate"
    else:
        return "high"

def _get_variable_usage(node: ast.AST) -> Set[str]:
    variables = set()
    for n in ast.walk(node):
        if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Store):
            variables.add(n.id)
    return variables

if __name__ == "__main__":
    mcp.run()