"""Enhanced MCP server for development with advanced code intelligence.

This MCP server provides comprehensive development tools:

1. Real-time file watching and change detection
2. Deep documentation research for any package (using LLM)
3. Semantic code search
4. Smart context building from multiple sources
5. Git integration with AI-powered diff reviews
6. Production-grade repository analysis
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    GetPromptResult,
    Prompt,
    PromptArgument,
    PromptMessage,
    TextContent,
    Tool,
)
from pydantic import BaseModel, Field

from .. import __version__ as KIT_VERSION
from ..ast_search import ASTSearcher
from ..pr_review.config import ReviewConfig
from ..pr_review.local_reviewer import LocalDiffReviewer
from ..repository import Repository

logger = logging.getLogger("kit-dev-mcp")

# MCP error codes
INVALID_PARAMS = -32602
INTERNAL_ERROR = -32603

# Context limits removed for better performance - let the AI handle context management


class MCPError(Exception):
    """MCP protocol error."""

    def __init__(self, code: int, message: str):
        self.code = code
        self.message = message
        super().__init__(message)


# Compatibility alias for different MCP versions
try:
    from mcp.types import EmbeddedResource

    ResourceContent: type[Union[EmbeddedResource, TextContent]] = EmbeddedResource
except ImportError:
    ResourceContent = TextContent


# Parameter classes from standard server
class OpenRepoParams(BaseModel):
    path_or_url: str
    github_token: Optional[str] = None
    ref: Optional[str] = None


class SearchParams(BaseModel):
    repo_id: str
    query: str
    pattern: str = "*.py"


class GrepParams(BaseModel):
    repo_id: str
    pattern: str
    case_sensitive: bool = True
    include_pattern: Optional[str] = None
    exclude_pattern: Optional[str] = None
    max_results: int = 1000
    directory: Optional[str] = None
    include_hidden: bool = False


class GetFileContentParams(BaseModel):
    repo_id: str
    file_path: Union[str, List[str]]


class GetMultipleFileContentsParams(BaseModel):
    repo_id: str
    file_paths: List[str]


class ExtractSymbolsParams(BaseModel):
    repo_id: str
    file_path: str
    symbol_type: Optional[str] = None


class FindSymbolUsagesParams(BaseModel):
    repo_id: str
    symbol_name: str
    symbol_type: Optional[str] = None
    file_path: Optional[str] = None


class GetFileTreeParams(BaseModel):
    repo_id: str


class GetCodeSummaryParams(BaseModel):
    repo_id: str
    file_path: str
    symbol_name: Optional[str] = None


class GitInfoParams(BaseModel):
    repo_id: str


class ReviewDiffParams(BaseModel):
    repo_id: str
    diff_spec: str
    priority_filter: Optional[List[str]] = None
    max_files: int = 10
    model: Optional[str] = None


class GrepASTParams(BaseModel):
    """Search code using AST patterns."""

    repo_id: str
    pattern: str = Field(description="AST pattern to search for")
    mode: str = Field(default="simple", description="Search mode: simple, pattern, or query")
    file_pattern: str = Field(default="**/*.py", description="File glob pattern")
    max_results: int = Field(default=20, ge=1, le=50)


class KitServerLogic:
    def __init__(self):
        self._repos: Dict[str, Repository] = {}

    def open_repository(self, path_or_url: str, github_token: Optional[str] = None, ref: Optional[str] = None) -> str:
        """Open a repository and return its ID."""
        import os

        # Check for GitHub token in environment if not provided
        if github_token is None:
            github_token = os.environ.get("KIT_GITHUB_TOKEN") or os.environ.get("GITHUB_TOKEN")

        try:
            repo = Repository(path_or_url, github_token=github_token, ref=ref)
            repo_id = f"repo_{int(time.time() * 1000)}"
            self._repos[repo_id] = repo
            return repo_id
        except Exception as e:
            raise MCPError(INVALID_PARAMS, f"Failed to open repository: {e}")

    def get_repo(self, repo_id: str) -> Repository:
        """Get a repository by ID."""
        if repo_id not in self._repos:
            raise MCPError(INVALID_PARAMS, f"Repository {repo_id} not found")
        return self._repos[repo_id]

    def get_file_content(self, repo_id: str, file_path: Union[str, List[str]]) -> Union[str, Dict[str, Any]]:
        """Get file content."""
        repo = self.get_repo(repo_id)
        try:
            if isinstance(file_path, list):
                result = {}
                for fp in file_path:
                    # Securely validate path to prevent traversal
                    try:
                        # Join the file path with repo path, then resolve
                        repo_path = Path(repo.repo_path).resolve()
                        safe_path = (repo_path / fp).resolve()

                        # Ensure resolved path is within repo bounds
                        if not safe_path.is_relative_to(repo_path):
                            raise MCPError(INVALID_PARAMS, f"Path traversal attempted: {fp}")

                        # Get the content directly with the original path
                        content = repo.get_file_content(fp)
                        result[fp] = content
                    except MCPError:
                        # Re-raise MCPError for path traversal
                        raise
                    except FileNotFoundError:
                        result[fp] = f"File not found: {fp}"
                    except ValueError as e:
                        if "outside repository" in str(e).lower():
                            raise MCPError(INVALID_PARAMS, f"Path traversal attempted: {e}")
                        raise MCPError(INVALID_PARAMS, str(e))
                    except Exception as e:
                        result[fp] = f"Error reading file: {e!s}"
                return result
            else:
                # Securely validate single file path
                repo_path = Path(repo.repo_path).resolve()
                safe_path = (repo_path / file_path).resolve()

                # Ensure resolved path is within repo bounds
                if not safe_path.is_relative_to(repo_path):
                    raise MCPError(INVALID_PARAMS, f"Path traversal attempted: {file_path}")

                # Get the content directly with the original path
                return repo.get_file_content(file_path)
        except ValueError as e:
            if "outside repository bounds" in str(e):
                raise MCPError(INVALID_PARAMS, f"Path traversal attempted: {e}")
            raise MCPError(INVALID_PARAMS, str(e))
        except FileNotFoundError as e:
            raise MCPError(INVALID_PARAMS, str(e))

    def get_multiple_file_contents(self, repo_id: str, file_paths: List[str]) -> Dict[str, Any]:
        """Get multiple file contents."""
        result = self.get_file_content(repo_id, file_paths)
        if isinstance(result, dict):
            return result
        # Should not happen as we're passing a list
        return {"error": "Unexpected result type"}

    def grep_code(
        self,
        repo_id: str,
        pattern: str,
        case_sensitive: bool = True,
        include_pattern: Optional[str] = None,
        exclude_pattern: Optional[str] = None,
        max_results: int = 1000,
        directory: Optional[str] = None,
        include_hidden: bool = False,
    ) -> List[Dict[str, Any]]:
        """Grep for patterns in code with context limits."""
        repo = self.get_repo(repo_id)
        try:
            results = repo.grep(
                pattern,
                case_sensitive=case_sensitive,
                include_pattern=include_pattern,
                exclude_pattern=exclude_pattern,
                max_results=max_results,
                directory=directory,
                include_hidden=include_hidden,
            )

            return results
        except ValueError as e:
            raise MCPError(INVALID_PARAMS, str(e))

    def get_file_tree(self, repo_id: str) -> List[Dict[str, Any]]:
        """Get file tree."""
        repo = self.get_repo(repo_id)
        return repo.get_file_tree()

    def extract_symbols(self, repo_id: str, file_path: str, symbol_type: Optional[str] = None) -> List[Dict[str, Any]]:
        """Extract symbols from a file."""
        repo = self.get_repo(repo_id)
        # Repository.extract_symbols doesn't accept symbol_type parameter
        try:
            return repo.extract_symbols(file_path)
        except ValueError as e:
            if "outside repository bounds" in str(e):
                raise MCPError(INVALID_PARAMS, f"Path traversal attempted: {e}")
            raise MCPError(INVALID_PARAMS, str(e))

    def find_symbol_usages(
        self, repo_id: str, symbol_name: str, symbol_type: Optional[str] = None, file_path: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """Find symbol usages."""
        repo = self.get_repo(repo_id)
        # Repository.find_symbol_usages doesn't accept keyword arguments
        return repo.find_symbol_usages(symbol_name)

    def get_code_summary(self, repo_id: str, file_path: str, symbol_name: Optional[str] = None) -> Dict[str, Any]:
        """Get code summary."""
        repo = self.get_repo(repo_id)
        # Repository doesn't have get_code_summary, so we'll implement a basic one
        try:
            content = repo.get_file_content(file_path)
            symbols = repo.extract_symbols(file_path)
        except (ValueError, FileNotFoundError) as e:
            raise MCPError(INVALID_PARAMS, str(e))

        summary = {
            "file": file_path,
            "content_preview": content[:500] if content else "",
            "symbols": symbols[:10] if symbols else [],  # First 10 symbols
            "line_count": len(content.splitlines()) if content else 0,
        }

        if symbol_name:
            # Find specific symbol
            matching_symbols = [s for s in symbols if s.get("name") == symbol_name]
            if matching_symbols:
                summary["requested_symbol"] = matching_symbols[0]

        return {"summary": summary}

    def get_git_info(self, repo_id: str) -> Dict[str, Any]:
        """Get git information."""
        repo = self.get_repo(repo_id)
        # Repository doesn't have get_git_info, construct it from available methods
        return {
            "current_sha": repo.current_sha,
            "current_sha_short": repo.current_sha_short,
            "current_branch": repo.current_branch,
            "remote_url": repo.remote_url,
            "is_dirty": repo.is_dirty,
            "tags": repo.tags,
            "branches": repo.branches,
        }

    def review_diff(
        self,
        repo_id: str,
        diff_spec: str,
        priority_filter: Optional[List[str]] = None,
        max_files: int = 10,
        model: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Review a diff."""

        repo = self.get_repo(repo_id)

        try:
            # Create review config
            config = ReviewConfig.from_file()
            if model:
                config.llm.model = model

            # Create reviewer
            reviewer = LocalDiffReviewer(config, repo.repo_path)

            # Perform review
            result = reviewer.review(diff_spec)

            # Extract cost from result if present
            import re

            cost_match = re.search(r"Cost: \$(\d+\.\d+)", result)
            cost = float(cost_match.group(1)) if cost_match else 0.0

            return {"review": result, "diff_spec": diff_spec, "cost": cost, "model": model or "gpt-4"}
        except Exception as e:
            raise MCPError(INTERNAL_ERROR, f"Failed to review diff: {e}")

    def search_code(self, repo_id: str, query: str, pattern: str = "*.py") -> List[Dict[str, Any]]:
        """Search code with context limits."""
        repo = self.get_repo(repo_id)
        # Repository uses search_text, not search_code
        results = repo.search_text(query, file_pattern=pattern)

        return results

    async def grep_ast_async(
        self, repo_id: str, pattern: str, mode: str = "simple", file_pattern: str = "**/*.py", max_results: int = 20
    ) -> List[Dict[str, Any]]:
        """Search code using AST patterns with async execution to prevent blocking.

        Examples:
            - pattern="async def" - Find all async functions
            - pattern="class $NAME(BaseModel)" - Find classes extending BaseModel
            - pattern='{"type": "try_statement"}' - Find all try blocks
        """
        repo = self.get_repo(repo_id)

        # Run the heavy AST parsing in a thread pool to avoid blocking the event loop
        loop = asyncio.get_event_loop()
        with ThreadPoolExecutor(max_workers=2) as executor:
            results = await loop.run_in_executor(
                executor, self._grep_ast_sync, repo.repo_path, pattern, file_pattern, mode, max_results
            )

        return results

    def _grep_ast_sync(
        self, repo_path: str, pattern: str, file_pattern: str, mode: str, max_results: int
    ) -> List[Dict[str, Any]]:
        """Synchronous AST search implementation for thread pool execution."""
        # Create AST searcher
        searcher = ASTSearcher(repo_path)

        # Perform search
        results = searcher.search_pattern(
            pattern=pattern, file_pattern=file_pattern, mode=mode, max_results=max_results
        )

        # Add a preview of the match
        for result in results:
            if "text" in result:
                lines = result["text"].split("\n")
                result["preview"] = lines[0] if lines else ""
                if len(result["preview"]) > 100:
                    result["preview"] = result["preview"][:100] + "..."

        return results

    def grep_ast(
        self, repo_id: str, pattern: str, mode: str = "simple", file_pattern: str = "**/*.py", max_results: int = 20
    ) -> List[Dict[str, Any]]:
        """Search code using AST patterns with context limits (synchronous wrapper).

        Examples:
            - pattern="async def" - Find all async functions
            - pattern="class $NAME(BaseModel)" - Find classes extending BaseModel
            - pattern='{"type": "try_statement"}' - Find all try blocks
        """
        # For backward compatibility, provide a sync wrapper
        # This is called from the MCP server, which will handle async execution
        repo = self.get_repo(repo_id)
        return self._grep_ast_sync(repo.repo_path, pattern, file_pattern, mode, max_results)

    def list_prompts(self) -> List[Any]:
        """List available prompts."""
        return [
            Prompt(
                name="review_diff",
                description="Review a local git diff",
                arguments=[
                    PromptArgument(name="repo_id", description="Repository ID", required=True),
                    PromptArgument(
                        name="diff_spec", description="Git diff spec (e.g., 'HEAD~1..HEAD', '--staged')", required=True
                    ),
                    PromptArgument(name="priority_filter", description="Priority levels to include", required=False),
                    PromptArgument(name="max_files", description="Maximum number of files to review", required=False),
                    PromptArgument(name="model", description="LLM model to use", required=False),
                ],
            )
        ]

    def get_prompt(self, name: str, arguments: dict) -> Any:
        """Get a specific prompt."""
        # This would normally handle prompt requests
        from mcp.types import GetPromptResult, PromptMessage, TextContent

        if name == "review_diff":
            repo_id = arguments.get("repo_id")
            diff_spec = arguments.get("diff_spec", "HEAD~1..HEAD")

            # Keep repo_id as string for review_diff
            if not repo_id:
                raise MCPError(INVALID_PARAMS, "repo_id is required")
            result = self.review_diff(repo_id, diff_spec)

            return GetPromptResult(
                description=f"AI review of diff: {diff_spec}",
                messages=[PromptMessage(role="assistant", content=TextContent(type="text", text=result["review"]))],
            )

        raise MCPError(INVALID_PARAMS, f"Unknown prompt: {name}")

    def list_tools(self) -> List[Tool]:
        """List standard Kit tools."""
        return [
            Tool(
                name="open_repository",
                description="Open a local or remote Git repository",
                inputSchema=OpenRepoParams.model_json_schema(),
            ),
            Tool(
                name="grep_code", description="Fast literal string search", inputSchema=GrepParams.model_json_schema()
            ),
            Tool(
                name="get_file_tree",
                description="Get repository file structure",
                inputSchema=GetFileTreeParams.model_json_schema(),
            ),
            Tool(
                name="extract_symbols",
                description="Extract symbols from a file",
                inputSchema=ExtractSymbolsParams.model_json_schema(),
            ),
            Tool(
                name="find_symbol_usages",
                description="Find where symbols are used",
                inputSchema=FindSymbolUsagesParams.model_json_schema(),
            ),
            Tool(
                name="review_diff",
                description="Review a local git diff with AI",
                inputSchema=ReviewDiffParams.model_json_schema(),
            ),
            Tool(
                name="grep_ast",
                description="Search code using AST patterns (semantic search)",
                inputSchema=GrepASTParams.model_json_schema(),
            ),
        ]


class BuildContextParams(BaseModel):
    """Build comprehensive context for a development task."""

    repo_id: str
    task_description: str
    include_tests: bool = True
    include_docs: bool = True
    include_dependencies: bool = True
    max_files: int = 20


class DeepResearchParams(BaseModel):
    """Deep research documentation for a package (using LLM)."""

    package_name: str
    query: str = Field(default="", description="Specific question about the package")


class LocalDevServerLogic(KitServerLogic):
    """Enhanced MCP server logic for development."""

    def __init__(self):
        super().__init__()
        self._test_results: Dict[str, Dict] = {}
        self._context_cache: Dict[str, Any] = {}

    def open_repository(self, path_or_url: str, github_token: Optional[str] = None, ref: Optional[str] = None) -> str:
        """Open a repository."""
        try:
            repo = Repository(path_or_url, github_token=github_token, ref=ref)
            repo_id = f"repo_{len(self._repos)}"
            self._repos[repo_id] = repo
            logger.info(f"Opened repository at {path_or_url} with ID {repo_id}")
            return repo_id
        except Exception as e:
            logger.error(f"Failed to open repository: {e}")
            raise

    def deep_research_package(self, package_name: str, query: str = "") -> Dict[str, Any]:
        """Perform deep research on a package using Context7 + LLM synthesis."""
        import json
        import os

        from ..context7_client import Context7Client, integrate_with_llm
        from ..deep_research import DeepResearch
        from ..summaries import AnthropicConfig, LLMError, OpenAIConfig

        # First, try to fetch real documentation from Context7
        context7 = Context7Client()
        context7_docs = context7.fetch_documentation(package_name, topic=query if query else None)

        # Check if we have LLM config for synthesis
        config: Optional[Union[OpenAIConfig, AnthropicConfig]] = None
        if os.environ.get("OPENAI_API_KEY"):
            config = OpenAIConfig(model="gpt-5")
        elif os.environ.get("ANTHROPIC_API_KEY"):
            config = AnthropicConfig(model="claude-3-haiku-20240307")

        # If Context7 returned good data and we have no LLM, return Context7 data
        if context7_docs.get("status") == "success" and not config:
            return {
                "package": package_name,
                "source": "context7",
                "model": "none",
                "execution_time": 0,
                "documentation": context7_docs.get("documentation", {}),
                "message": "Using Context7 aggregated documentation (no LLM synthesis)",
            }

        # If no LLM config and Context7 failed, return helpful message
        if not config:
            return {
                "package": package_name,
                "source": "none",
                "model": "none",
                "execution_time": 0,
                "documentation": "No documentation sources available. To get comprehensive docs, either set OPENAI_API_KEY/ANTHROPIC_API_KEY for LLM synthesis, or ensure Context7 is accessible.",
            }

        try:
            researcher = DeepResearch(config)
        except LLMError as e:
            # If LLM fails but Context7 worked, return Context7 data
            if context7_docs.get("status") == "success":
                return {
                    "package": package_name,
                    "source": "context7",
                    "model": "error",
                    "execution_time": 0,
                    "documentation": context7_docs.get("documentation", {}),
                    "error": f"LLM error: {e}, using Context7 data only",
                }
            return {
                "package": package_name,
                "source": "error",
                "model": "error",
                "execution_time": 0,
                "documentation": f"LLM configuration error: {e}",
            }

        # Build the research query, including Context7 data if available
        context_info = ""
        if context7_docs.get("status") == "success":
            context_info = f"\n\nContext7 provided this real documentation:\n{json.dumps(context7_docs.get('documentation', {}), indent=2)[:2000]}\n\nPlease synthesize and enhance this information."

        if query:
            research_query = f"Tell me about the {package_name} package, specifically: {query}{context_info}"
        else:
            research_query = f"""Provide comprehensive information about the {package_name} Python package including:
            - What it does and its main use cases
            - Key features and capabilities
            - Basic usage examples
            - Common patterns and best practices
            - Important classes/functions to know about{context_info}"""

        result = researcher.research(research_query)

        # Integrate Context7 docs with LLM response
        if context7_docs.get("status") == "success":
            final_result = integrate_with_llm(context7_docs, result.answer)
            return {
                "package": package_name,
                "source": "context7+llm",
                "model": result.model,
                "execution_time": result.execution_time,
                "documentation": final_result,
            }

        # Just LLM response if Context7 didn't work
        return {
            "package": package_name,
            "source": "llm_only",
            "model": result.model,
            "execution_time": result.execution_time,
            "documentation": result.answer,
        }

    def build_smart_context(
        self,
        repo_id: str,
        task_description: str,
        include_tests: bool,
        include_docs: bool,
        include_dependencies: bool,
        max_files: int,
    ) -> Dict[str, Any]:
        """Build comprehensive context for a development task."""
        repo = self.get_repo(repo_id)

        context: Dict[str, Any] = {
            "task": task_description,
            "repository": {"path": repo.repo_path, "branch": repo.current_branch, "sha": repo.current_sha},
            "relevant_files": [],
            "symbols": [],
            "tests": [],
            "documentation": [],
            "dependencies": [],
        }

        keywords = task_description.lower().split()

        for file_info in repo.get_file_tree()[:max_files]:
            if file_info.get("is_dir"):
                continue

            file_path = file_info["path"]

            # Check if file might be relevant
            if any(keyword in file_path.lower() for keyword in keywords):
                context["relevant_files"].append({"path": file_path, "size": file_info.get("size", 0)})

                # Extract symbols from relevant files
                try:
                    symbols = repo.extract_symbols(file_path)
                    context["symbols"].extend(
                        [
                            {"file": file_path, "name": s["name"], "type": s["type"]}
                            for s in symbols[:5]  # Limit symbols per file
                        ]
                    )
                except Exception:
                    pass

        # Find tests if requested
        if include_tests:
            test_patterns = ["test_*.py", "*_test.py", "*.test.js", "*.spec.js"]
            for file_info in repo.get_file_tree():
                if any(file_info["path"].endswith(pattern.replace("*", "")) for pattern in test_patterns):
                    context["tests"].append(file_info["path"])

        # Get dependencies if requested
        if include_dependencies:
            requirements_file = Path(repo.repo_path) / "requirements.txt"
            if requirements_file.exists():
                try:
                    deps = []
                    with open(requirements_file, "r") as f:
                        for line in f:
                            line = line.strip()
                            if line and not line.startswith("#"):
                                # Simple parsing of requirement lines
                                if "==" in line:
                                    name, version = line.split("==")
                                    deps.append({"name": name.strip(), "version": version.strip()})
                                elif ">=" in line:
                                    name = line.split(">=")[0]
                                    deps.append({"name": name.strip(), "version": "latest"})
                                else:
                                    deps.append({"name": line, "version": "any"})
                    context["dependencies"] = deps[:10]
                except Exception as e:
                    logger.warning(f"Failed to parse requirements.txt: {e}")
                    context["dependencies"] = []
            else:
                context["dependencies"] = []

        # Cache context
        cache_key = f"{repo_id}:{task_description[:50]}"
        self._context_cache[cache_key] = context

        return context

    def list_tools(self) -> List[Tool]:
        """List all available tools."""
        # Get base tools from parent class
        parent_tools = super().list_tools()

        # Add our enhanced development tools
        dev_tools = [
            # Documentation and research
            Tool(
                name="deep_research_package",
                description="Deep research documentation for any package (using LLM to provide comprehensive answers)",
                inputSchema=DeepResearchParams.model_json_schema(),
            ),
            # Smart context building
            Tool(
                name="build_smart_context",
                description="Build comprehensive context for a development task",
                inputSchema=BuildContextParams.model_json_schema(),
            ),
        ]

        # Return combined tools
        return parent_tools + dev_tools


async def serve():
    """Serve the enhanced development MCP server."""
    server = Server("kit-dev-mcp", version=KIT_VERSION)
    logic = LocalDevServerLogic()

    @server.call_tool()
    async def call_tool(name: str, arguments: dict) -> List[TextContent]:
        """Handle tool calls."""
        try:
            # Handle development-specific tools
            if name == "deep_research_package":
                research_params = DeepResearchParams(**arguments)
                result = logic.deep_research_package(research_params.package_name, research_params.query)
                return [TextContent(type="text", text=json.dumps(result, indent=2))]

            elif name == "build_smart_context":
                build_params = BuildContextParams(**arguments)
                result = logic.build_smart_context(
                    build_params.repo_id,
                    build_params.task_description,
                    build_params.include_tests,
                    build_params.include_docs,
                    build_params.include_dependencies,
                    build_params.max_files,
                )
                return [TextContent(type="text", text=json.dumps(result, indent=2))]

            # For all other tools, delegate to parent class handling
            elif name in [
                "open_repository",
                "search_code",
                "grep_code",
                "get_file_content",
                "get_multiple_file_contents",
                "extract_symbols",
                "find_symbol_usages",
                "get_file_tree",
                "get_code_summary",
                "get_git_info",
                "review_diff",
            ]:
                # Route to parent class method
                if name == "open_repository":
                    open_params = OpenRepoParams(**arguments)
                    repo_id = logic.open_repository(open_params.path_or_url, open_params.github_token, open_params.ref)
                    return [TextContent(type="text", text=f"Opened repository with ID: {repo_id}")]
                elif name == "search_code":
                    search_params = SearchParams(**arguments)
                    result = logic.search_code(search_params.repo_id, search_params.query, search_params.pattern)
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "grep_code":
                    grep_params = GrepParams(**arguments)
                    result = logic.grep_code(
                        grep_params.repo_id,
                        grep_params.pattern,
                        grep_params.case_sensitive,
                        grep_params.include_pattern,
                        grep_params.exclude_pattern,
                        grep_params.max_results,
                        grep_params.directory,
                        grep_params.include_hidden,
                    )
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "get_file_content":
                    file_params = GetFileContentParams(**arguments)
                    result = logic.get_file_content(file_params.repo_id, file_params.file_path)
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "get_multiple_file_contents":
                    multi_params = GetMultipleFileContentsParams(**arguments)
                    result = logic.get_multiple_file_contents(multi_params.repo_id, multi_params.file_paths)
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "extract_symbols":
                    symbol_params = ExtractSymbolsParams(**arguments)
                    result = logic.extract_symbols(
                        symbol_params.repo_id, symbol_params.file_path, symbol_params.symbol_type
                    )
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "find_symbol_usages":
                    usage_params = FindSymbolUsagesParams(**arguments)
                    result = logic.find_symbol_usages(
                        usage_params.repo_id, usage_params.symbol_name, usage_params.file_path
                    )
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "get_file_tree":
                    tree_params = GetFileTreeParams(**arguments)
                    result = logic.get_file_tree(tree_params.repo_id)
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "get_code_summary":
                    summary_params = GetCodeSummaryParams(**arguments)
                    result = logic.get_code_summary(
                        summary_params.repo_id, summary_params.file_path, summary_params.symbol_name
                    )
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "get_git_info":
                    git_params = GitInfoParams(**arguments)
                    result = logic.get_git_info(git_params.repo_id)
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "review_diff":
                    review_params = ReviewDiffParams(**arguments)
                    result = logic.review_diff(
                        review_params.repo_id,
                        review_params.diff_spec,
                        review_params.priority_filter,
                        review_params.max_files,
                        review_params.model,
                    )
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                elif name == "grep_ast":
                    ast_params = GrepASTParams(**arguments)
                    result = logic.grep_ast(
                        ast_params.repo_id,
                        ast_params.pattern,
                        ast_params.mode,
                        ast_params.file_pattern,
                        ast_params.max_results,
                    )
                    return [TextContent(type="text", text=json.dumps(result, indent=2))]
                else:
                    # Should not happen since we checked the name is in the list
                    return [TextContent(type="text", text=f"Tool {name} is recognized but not implemented")]

            else:
                return [TextContent(type="text", text=f"Unknown tool: {name}")]

        except Exception as e:
            logger.error(f"Error in tool {name}: {e}")
            return [TextContent(type="text", text=f"Error: {e!s}")]

    @server.list_tools()
    async def list_tools() -> List[Tool]:
        """List available tools."""
        return logic.list_tools()

    @server.list_prompts()
    async def list_prompts() -> List[Prompt]:
        """List available prompts."""
        return [
            Prompt(
                name="analyze_codebase",
                description="Comprehensive codebase analysis with all features",
                arguments=[
                    PromptArgument(name="path", description="Path to the repository", required=True),
                    PromptArgument(name="task", description="What you want to accomplish", required=True),
                ],
            ),
        ]

    @server.get_prompt()
    async def get_prompt(name: str, arguments: dict) -> GetPromptResult:
        """Get a specific prompt."""
        if name == "analyze_codebase":
            return GetPromptResult(
                messages=[
                    PromptMessage(
                        role="user",
                        content=TextContent(
                            type="text",
                            text=f"Analyze the codebase at {arguments['path']} "
                            f"for the task: {arguments['task']}. "
                            "Use all available tools including file watching, "
                            "test running, performance analysis, and context building.",
                        ),
                    )
                ]
            )
        else:
            raise MCPError(INVALID_PARAMS, f"Unknown prompt: {name}")

    options = server.create_initialization_options()

    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, options)


if __name__ == "__main__":
    import asyncio

    asyncio.run(serve())
