server.py

"""
server.py
---------
CodeMap MCP Server — entry point.
Registers all six MCP tools and starts the server.

Usage:
    python server.py
    # or via MCP client config pointing to this file

Environment variables:
    ANTHROPIC_API_KEY  — required for generate_roadmap LLM output
                         (structural output works without it)
"""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# MCP SDK import
# ---------------------------------------------------------------------------
try:
    from mcp.server import Server
    from mcp.server.stdio import stdio_server
    from mcp import types
    HAS_MCP = True
except ImportError:
    HAS_MCP = False
    print(
        "ERROR: mcp package not installed.\n"
        "Install with: pip install mcp\n",
        file=sys.stderr,
    )
    sys.exit(1)

# ---------------------------------------------------------------------------
# CodeMap internal imports
# ---------------------------------------------------------------------------
from config import load_config, describe_config
from tools.discover import discover_structure
from tools.read_file import read_file
from tools.search_symbol import search_symbol
from tools.generate_roadmap import generate_roadmap
from tools.explain_error import explain_error
from tools.explain_pr import explain_pr

# ---------------------------------------------------------------------------
# Server setup
# ---------------------------------------------------------------------------
server = Server("codemap-mcp")

# ---------------------------------------------------------------------------
# Tool: discover_structure
# ---------------------------------------------------------------------------
@server.list_tools()
async def list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name="discover_structure",
            description=(
                "Walk the target repository or folder and return a structured map "
                "of all Python files, packages, and top-level directories. "
                "Read-only — uses os.walk and pathlib only."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "repo_path": {
                        "type": "string",
                        "description": "Absolute path to the repository root.",
                    },
                    "target": {
                        "type": "string",
                        "description": (
                            "Optional subfolder or feature name to narrow the scan. "
                            "Leave empty to scan the whole repository."
                        ),
                    },
                },
                "required": ["repo_path"],
            },
        ),
        types.Tool(
            name="read_file",
            description=(
                "Return the raw content of a single Python file so the LLM can "
                "inspect it in detail. Read-only — uses pathlib.read_text() only. "
                "Validates that the path stays within the repository root."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "repo_path": {
                        "type": "string",
                        "description": "Absolute path to the repository root.",
                    },
                    "file_path": {
                        "type": "string",
                        "description": "Path to file relative to repo_path.",
                    },
                },
                "required": ["repo_path", "file_path"],
            },
        ),
        types.Tool(
            name="search_symbol",
            description=(
                "Search for a class name, function name, or import pattern across "
                "all Python files in scope. Uses ast.parse() — no code execution. "
                "Returns file path, line number, and symbol type for each match."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "repo_path": {
                        "type": "string",
                        "description": "Absolute path to the repository root.",
                    },
                    "query": {
                        "type": "string",
                        "description": "Name or pattern to find (case-insensitive substring match).",
                    },
                    "scope": {
                        "type": "string",
                        "description": "Optional subfolder to limit search.",
                    },
                },
                "required": ["repo_path", "query"],
            },
        ),
        types.Tool(
            name="generate_roadmap",
            description=(
                "Use the discovered structure and symbol graph to generate a "
                "human-readable roadmap. Produces prose flow explanation, key files "
                "table, and a Mermaid graph with conditional edge labels. "
                "Requires ANTHROPIC_API_KEY for LLM prose; structural output available without it."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "repo_path": {
                        "type": "string",
                        "description": "Absolute path to the repository root.",
                    },
                    "scope": {
                        "type": "string",
                        "description": (
                            "Folder path or feature name to explain "
                            "(e.g. 'payments', 'payments/', 'app/services')."
                        ),
                    },
                    "question": {
                        "type": "string",
                        "description": (
                            "Optional specific aspect to explain "
                            "(e.g. 'how is process_payment triggered?')."
                        ),
                    },
                    "max_depth": {
                        "type": "integer",
                        "description": "Maximum call chain nodes to trace (default 15, range 1-20).",
                        "minimum": 1,
                        "maximum": 20,
                        "default": 15,
                    },
                    "use_llm": {
                        "type": "boolean",
                        "description": (
                            "Set to false to skip LLM and return structural output only. "
                            "Default: true."
                        ),
                        "default": True,
                    },
                },
                "required": ["repo_path", "scope"],
            },
        ),
        types.Tool(
            name="explain_error",
            description=(
                "Parse a Python stacktrace/traceback and explain the issue with local "
                "source code context. Uses static search (read-only)."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "repo_path": {
                        "type": "string",
                        "description": "Absolute path to the repository root.",
                    },
                    "stacktrace": {
                        "type": "string",
                        "description": "The pasted stacktrace/traceback string.",
                    },
                    "use_llm": {
                        "type": "boolean",
                        "description": "Set to false to skip LLM and return structural context only. Default: true.",
                        "default": True,
                    },
                },
                "required": ["repo_path", "stacktrace"],
            },
        ),
        types.Tool(
            name="explain_pr",
            description=(
                "Analyse the impact of a pull request or git diff without pushing data "
                "to any external service. Returns changed files, changed symbols, blast "
                "radius, and affected HTTP routes as Markdown. Read-only — uses "
                "git diff + static call graph only."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "repo_path": {
                        "type": "string",
                        "description": "Absolute path to the repository root.",
                    },
                    "base_ref": {
                        "type": "string",
                        "description": "Git ref (branch or commit SHA) to diff against (default: 'main').",
                        "default": "main",
                    },
                    "head_ref": {
                        "type": "string",
                        "description": "Git ref for the PR head (default: 'HEAD').",
                        "default": "HEAD",
                    },
                    "diff_text": {
                        "type": "string",
                        "description": (
                            "Optional: paste a raw unified diff or list of changed files here. "
                            "When provided, skips running git and uses this text directly."
                        ),
                    },
                    "use_llm": {
                        "type": "boolean",
                        "description": (
                            "Set to false to skip the AI impact summary and return "
                            "structural analysis only. Default: true."
                        ),
                        "default": True,
                    },
                },
                "required": ["repo_path"],
            },
        ),
    ]


# ---------------------------------------------------------------------------
# Tool call dispatcher
# ---------------------------------------------------------------------------
async def _run_tool_handler(
    name: str,
    arguments: dict[str, Any],
) -> list[types.TextContent]:
    """
    Dispatch MCP tool calls to the appropriate handler.
    Always returns a list of TextContent — never writes to disk.
    """
    repo_path = arguments.get("repo_path", "")

    # Load config for this repo (cached per call for simplicity)
    config = {}
    if repo_path:
        try:
            config = load_config(repo_path)
        except Exception:
            config = {}

    # -----------------------------------------------------------------------
    # discover_structure
    # -----------------------------------------------------------------------
    if name == "discover_structure":
        result = discover_structure(
            repo_path=repo_path,
            target=arguments.get("target"),
            config=config,
        )
        if result["success"]:
            output = result["tree_text"]
            stats = result["file_info"].get("stats", {})
            output += (
                f"\n\nStats: {stats.get('total_py_files', 0)} Python files, "
                f"{stats.get('total_packages', 0)} packages, "
                f"{stats.get('skipped', 0)} skipped."
            )
        else:
            output = f"Error: {result['error']}"
        return [types.TextContent(type="text", text=output)]

    # -----------------------------------------------------------------------
    # read_file
    # -----------------------------------------------------------------------
    elif name == "read_file":
        result = read_file(
            repo_path=repo_path,
            file_path=arguments.get("file_path", ""),
        )
        if result["success"]:
            output = (
                f"File: {result['path']}  ({result['lines']} lines, {result['size_kb']} KB)\n"
                f"{'=' * 60}\n"
                f"{result['content']}"
            )
        else:
            output = f"Error: {result['error']}"
        return [types.TextContent(type="text", text=output)]

    # -----------------------------------------------------------------------
    # search_symbol
    # -----------------------------------------------------------------------
    elif name == "search_symbol":
        result = search_symbol(
            repo_path=repo_path,
            query=arguments.get("query", ""),
            scope=arguments.get("scope"),
            config=config,
        )
        if result["success"]:
            matches = result["matches"]
            if not matches:
                output = f"No symbols matching '{arguments.get('query')}' found."
            else:
                lines = [
                    f"Found {result['total_matches']} match(es) "
                    f"across {result['files_scanned']} files:\n"
                ]
                for m in matches:
                    parent = f" (in {m['parent']})" if m.get("parent") else ""
                    doc = f"\n    {m['docstring']}" if m.get("docstring") else ""
                    lines.append(
                        f"  [{m['kind']}] {m['name']}{parent}\n"
                        f"    {m['file']}:{m['line']}{doc}"
                    )
                if result["parse_errors"]:
                    lines.append(
                        f"\n⚠️  {len(result['parse_errors'])} file(s) had parse errors and were skipped."
                    )
                output = "\n".join(lines)
        else:
            output = f"Error: {result['error']}"
        return [types.TextContent(type="text", text=output)]

    # -----------------------------------------------------------------------
    # generate_roadmap
    # -----------------------------------------------------------------------
    elif name == "generate_roadmap":
        max_depth = arguments.get("max_depth", 15)
        # Clamp to valid range
        max_depth = max(1, min(20, int(max_depth)))

        result = generate_roadmap(
            repo_path=repo_path,
            scope=arguments.get("scope", ""),
            question=arguments.get("question", ""),
            max_depth=max_depth,
            config=config,
            use_llm=arguments.get("use_llm", True),
        )
        if result["success"]:
            output_parts = [result["roadmap"]]
            if result["routes_found"]:
                output_parts.append(
                    f"\n_HTTP routes found: {result['routes_found']}_"
                )
        else:
            output_parts = [f"Error generating roadmap: {result['error']}"]

        output = "\n".join(output_parts)
        return [types.TextContent(type="text", text=output)]

    # -----------------------------------------------------------------------
    # explain_error
    # -----------------------------------------------------------------------
    elif name == "explain_error":
        result = explain_error(
            repo_path=repo_path,
            stacktrace=arguments.get("stacktrace", ""),
            use_llm=arguments.get("use_llm", True),
        )
        if result["success"]:
            output = result["explanation"]
        else:
            output = f"Error explaining stacktrace: {result['error']}"
        return [types.TextContent(type="text", text=output)]

    # -----------------------------------------------------------------------
    # explain_pr
    # -----------------------------------------------------------------------
    elif name == "explain_pr":
        result = explain_pr(
            repo_path=repo_path,
            base_ref=arguments.get("base_ref", "main"),
            head_ref=arguments.get("head_ref", "HEAD"),
            diff_text=arguments.get("diff_text", ""),
            use_llm=arguments.get("use_llm", True),
            config=config,
        )
        if result["success"]:
            output = result["summary_md"]
        else:
            output = f"Error analysing PR diff: {result['error']}"
        return [types.TextContent(type="text", text=output)]

    # -----------------------------------------------------------------------
    # Unknown tool
    # -----------------------------------------------------------------------
    else:
        return [
            types.TextContent(
                type="text",
                text=f"Unknown tool: '{name}'. Available tools: "
                     "discover_structure, read_file, search_symbol, "
                     "generate_roadmap, explain_error, explain_pr",
            )
        ]


@server.call_tool()
async def call_tool(
    name: str,
    arguments: dict[str, Any],
) -> list[types.TextContent]:
    """
    Dispatch MCP tool calls to the appropriate handler.
    Supports local paths and remote Git URLs (automatically clones and cleans up).
    """
    repo_path = arguments.get("repo_path", "")
    is_remote = repo_path.startswith(("http://", "https://", "git@", "git://"))

    if is_remote:
        from analysis.remote_repo import clone_remote_repo
        cwd = Path(".").resolve()
        try:
            with clone_remote_repo(repo_path, cwd) as temp_local_path:
                local_args = dict(arguments)
                local_args["repo_path"] = str(temp_local_path)
                return await _run_tool_handler(name, local_args)
        except Exception as e:
            return [
                types.TextContent(
                    type="text",
                    text=f"Error executing tool '{name}' on remote repository: {e}"
                )
            ]
    else:
        return await _run_tool_handler(name, arguments)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
async def main():
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            server.create_initialization_options(),
        )


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())