Azure AI Search MCP Server

"""MCP Server implementation for Azure AI Search."""

import asyncio
import logging
import sys
import time
from typing import Any, Dict, List, Optional


try:
    from mcp.server import Server
    from mcp.server.stdio import stdio_server
    from mcp.types import (
        GetPromptResult,
        Prompt,
        PromptArgument,
        PromptMessage,
        Resource,
        TextContent,
        Tool,
    )
    from pydantic import AnyUrl

    MCP_AVAILABLE = True
except ImportError:
    # Use stderr to avoid interfering with JSON-RPC stdout
    print(
        "MCP library not available. Please install with: pip install mcp",
        file=sys.stderr,
    )
    MCP_AVAILABLE = False

from langsmith import traceable

from .chain import AzureSearchChain
from .config import config
from .prompt_manager import PromptManager


# Set up logging to stderr to avoid interfering with JSON-RPC stdout
logging.basicConfig(
    level=getattr(logging, config.mcp_server.log_level),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stderr)],
)
logger = logging.getLogger(__name__)


class AzureSearchMCPServer:
    """Azure AI Search MCP Server."""

    def __init__(self):
        """Initialize the MCP server."""
        self._init_warnings = []

        # Check configuration before initializing components
        self._check_configuration()

        # Initialize components (even if config is incomplete - they'll handle it gracefully)
        self.search_chain = AzureSearchChain()
        self.prompt_manager = PromptManager()

        if MCP_AVAILABLE:
            self.server = Server(config.mcp_server.name)
            self._setup_handlers()

    def _check_configuration(self):
        """Check and report configuration issues."""
        missing_azure = []
        if not config.azure_search.endpoint:
            missing_azure.append("AZURE_SEARCH_ENDPOINT")
        if not config.azure_search.api_key:
            missing_azure.append("AZURE_SEARCH_API_KEY")
        if not config.azure_search.index_name:
            missing_azure.append("AZURE_SEARCH_INDEX_NAME")

        if missing_azure:
            warning = f"Azure Search configuration incomplete. Missing: {', '.join(missing_azure)}. Search functionality will be disabled."
            self._init_warnings.append(warning)
            print(f"[CONFIG_WARNING] {warning}", file=sys.stderr)

        if not config.gemini.api_key:
            warning = "Gemini API key not found. LLM-based formatting will be disabled."
            self._init_warnings.append(warning)
            print(f"[CONFIG_WARNING] {warning}", file=sys.stderr)

    def _setup_handlers(self):
        """Set up MCP server handlers."""
        if not MCP_AVAILABLE:
            return

        @self.server.list_resources()
        async def handle_list_resources() -> List[Resource]:
            """List available resources."""

            resources = []

            # Add each tool as a resource
            tools = self._get_available_tools()
            for tool in tools:
                resources.append(
                    Resource(
                        uri=AnyUrl(f"tool://{tool.name}"),
                        name=tool.name,
                        description=tool.description,
                        mimeType="application/json",
                    )
                )

            # Add each prompt as a resource
            prompts = self._get_available_prompts()
            for prompt in prompts:
                resources.append(
                    Resource(
                        uri=AnyUrl(f"prompt://{prompt.name}"),
                        name=prompt.name,
                        description=prompt.description,
                        mimeType="text/plain",
                    )
                )

            return resources

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

        @self.server.list_prompts()
        async def handle_list_prompts() -> List[Prompt]:
            """List available prompts."""
            return self._get_available_prompts()

        @self.server.get_prompt()
        async def handle_get_prompt(name: str, arguments: Dict[str, str]) -> GetPromptResult:
            """Get a specific prompt."""
            try:
                if arguments is None:
                    arguments = {}

                query = arguments.get("query", "")
                documents = arguments.get("documents", "")

                if name.startswith("format_"):
                    # Extract format type
                    format_type = name[7:]  # Remove "format_" prefix

                    # Get the prompt template for this format
                    template = self.prompt_manager.get_prompt_template_for_format(
                        format_type
                    )

                    # Format the prompt with the provided arguments
                    formatted_prompt = template.format_messages(
                        query=query, documents=documents
                    )

                    return GetPromptResult(
                        description=f"Format documents in {format_type} style",
                        messages=[
                            PromptMessage(
                                role="user",
                                content=TextContent(
                                    type="text",
                                    text=str(msg.content)
                                    if hasattr(msg, "content")
                                    else str(msg),
                                ),
                            )
                            for msg in formatted_prompt
                        ],
                    )

                elif name.startswith("persona_"):
                    # Extract persona type
                    persona_type = name[8:]  # Remove "persona_" prefix

                    # Get persona information
                    persona_info = self.prompt_manager.get_persona_info(persona_type)

                    # Create basic template
                    basic_template = f"""Acting as {persona_info["name"]}: {persona_info["description"]}

Please respond to the following query with the characteristics and expertise of {persona_info["name"]}:

<query>
{query}
</query>

**Response:**"""

                    return GetPromptResult(
                        description=f"Persona prompt for {persona_info['name']}",
                        messages=[
                            PromptMessage(
                                role="user",
                                content=TextContent(type="text", text=basic_template),
                            )
                        ],
                    )

                else:
                    error_text = f"Unknown prompt: {name}"
                    return GetPromptResult(
                        description="Error",
                        messages=[
                            PromptMessage(
                                role="user",
                                content=TextContent(type="text", text=error_text),
                            )
                        ],
                    )

            except Exception as e:
                logger.error(f"Error getting prompt {name}: {str(e)}")
                error_text = f"Error getting prompt: {str(e)}"
                return GetPromptResult(
                    description="Error",
                    messages=[
                        PromptMessage(
                            role="user",
                            content=TextContent(type="text", text=error_text),
                        )
                    ],
                )

        @self.server.call_tool()
        @traceable(name="mcp_tool_call")
        async def handle_call_tool(
            name: str, arguments: Dict[str, Any]
        ) -> List[TextContent]:
            """Handle tool calls."""
            import time

            start_time = time.time()

            # Add a timeout wrapper to prevent hanging
            async def _execute_tool():
                try:
                    if name == "search_documents":
                        # Check if configuration is valid before attempting search
                        if self._init_warnings:
                            warning_msg = (
                                "Search functionality is disabled due to configuration issues:\n"
                                + "\n".join(f"• {w}" for w in self._init_warnings)
                            )
                            return [
                                TextContent(
                                    type="text",
                                    text=f"❌ {warning_msg}\n\nPlease set the required environment variables and restart the server.",
                                )
                            ]

                        query = arguments["query"]
                        top_k = arguments.get("top_k", 5)
                        format_type = arguments.get("format_type", "structured")

                        search_result = await self.search_chain.run(
                            query, top_k=top_k, output_format=format_type
                        )

                        result = search_result["context"]
                        text_content = TextContent(type="text", text=result)
                        return [text_content]

                    elif name == "search_and_summarize":
                        # Check if configuration is valid before attempting search
                        if self._init_warnings:
                            warning_msg = (
                                "Search functionality is disabled due to configuration issues:\n"
                                + "\n".join(f"• {w}" for w in self._init_warnings)
                            )
                            return [
                                TextContent(
                                    type="text",
                                    text=f"❌ {warning_msg}\n\nPlease set the required environment variables and restart the server.",
                                )
                            ]

                        query = arguments["query"]
                        top_k = arguments.get("top_k", 5)

                        search_result = await self.search_chain.run(
                            query, top_k=top_k, output_format="summary"
                        )

                        result = search_result["context"]
                        return [TextContent(type="text", text=result)]

                    elif name == "search_with_analysis":
                        # Check if configuration is valid before attempting search
                        if self._init_warnings:
                            warning_msg = (
                                "Search functionality is disabled due to configuration issues:\n"
                                + "\n".join(f"• {w}" for w in self._init_warnings)
                            )
                            return [
                                TextContent(
                                    type="text",
                                    text=f"❌ {warning_msg}\n\nPlease set the required environment variables and restart the server.",
                                )
                            ]

                        query = arguments["query"]
                        top_k = arguments.get("top_k", 5)

                        search_result = await self.search_chain.run(
                            query, top_k=top_k, output_format="analysis"
                        )

                        result = search_result["context"]
                        return [TextContent(type="text", text=result)]

                    elif name == "get_document_context":
                        document_ids = arguments["document_ids"]

                        result = await self.search_chain.get_document_context_tool(
                            document_ids
                        )
                        return [TextContent(type="text", text=result)]

                    else:
                        error_msg = f"Unknown tool: {name}"
                        logger.error(error_msg)
                        return [TextContent(type="text", text=f"Error: {error_msg}")]

                except asyncio.CancelledError:
                    elapsed = time.time() - start_time
                    logger.warning(f"Tool call cancelled: {name} after {elapsed:.2f}s")
                    raise  # Re-raise to properly handle cancellation

                except Exception as e:
                    print(f"!!! EXCEPTION MESSAGE: {str(e)} !!!", file=sys.stderr)
                    logger.error(f"Critical error in tool {name}: {e}")
                    raise

            # Run with timeout
            try:
                return await asyncio.wait_for(_execute_tool(), timeout=30.0)
            except asyncio.TimeoutError:
                elapsed = time.time() - start_time
                logger.error(f"Tool call timeout: {name} after {elapsed:.2f}s")
                return [
                    TextContent(
                        type="text",
                        text=f"Error: Tool call timed out after {elapsed:.2f}s. The search service may be experiencing delays.",
                    )
                ]
            except Exception as e:
                elapsed = time.time() - start_time
                logger.error(f"Tool call error: {name} after {elapsed:.2f}s: {e}")
                return [
                    TextContent(type="text", text=f"Error: Tool call failed: {str(e)}")
                ]

    def _get_available_tools(self) -> List[Tool]:
        """Get list of available tools."""
        return [
            Tool(
                name="search_documents",
                description="Search for relevant documents in Azure AI Search with AI-enhanced analysis and formatting",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query to find relevant documents",
                        },
                        "top_k": {
                            "type": "integer",
                            "description": "Number of documents to return (default: 5)",
                            "default": 5,
                        },
                        "format_type": {
                            "type": "string",
                            "description": "Format type: 'summary' (AI-enhanced default), 'structured', or 'analysis'",
                            "enum": ["summary", "structured", "analysis"],
                            "default": "summary",
                        },
                    },
                    "required": ["query"],
                },
            ),
            Tool(
                name="search_and_summarize",
                description="Search for documents and provide a concise summary of findings",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query to find relevant documents",
                        },
                        "top_k": {
                            "type": "integer",
                            "description": "Number of documents to return (default: 5)",
                            "default": 5,
                        },
                    },
                    "required": ["query"],
                },
            ),
            Tool(
                name="search_with_analysis",
                description="Search for documents and provide detailed relevance analysis",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query to find relevant documents",
                        },
                        "top_k": {
                            "type": "integer",
                            "description": "Number of documents to return (default: 5)",
                            "default": 5,
                        },
                    },
                    "required": ["query"],
                },
            ),
            Tool(
                name="get_document_context",
                description="Retrieve specific documents by their IDs for detailed context",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "document_ids": {
                            "type": "string",
                            "description": "Comma-separated list of document IDs to retrieve",
                        },
                        "format_output": {
                            "type": "boolean",
                            "description": "Whether to format the output for readability (default: true)",
                            "default": True,
                        },
                    },
                    "required": ["document_ids"],
                },
            ),
        ]

    def _get_available_prompts(self) -> List[Prompt]:
        """Get list of available prompts."""
        return [
            Prompt(
                name="search_query_optimization",
                description="Optimize search queries for better results",
                arguments=[
                    PromptArgument(
                        name="query",
                        description="The search query to optimize",
                        required=True,
                    )
                ],
            ),
            Prompt(
                name="document_analysis",
                description="Analyze document relevance and extract insights",
                arguments=[
                    PromptArgument(
                        name="documents",
                        description="Documents to analyze",
                        required=True,
                    ),
                    PromptArgument(
                        name="query",
                        description="The original search query",
                        required=True,
                    ),
                ],
            ),
        ]

    async def run_standalone_mode(self):
        """Run in standalone mode for testing without MCP."""
        logger.info("Running in standalone mode (MCP not available)")

        # Simple command-line interface for testing
        print("Azure AI Search Server - Standalone Mode")
        print("Available commands:")
        print("  search <query> - Search for documents (structured format)")
        print("  summary <query> - Search and provide summary")
        print("  analysis <query> - Search with relevance analysis")
        print("  get <doc_ids> - Get document context")
        print("  quit - Exit")

        while True:
            try:
                command = input("\n> ").strip()
                if not command:
                    continue

                if command.lower() in ["quit", "exit"]:
                    break

                parts = command.split(" ", 1)
                if len(parts) < 2:
                    print("Please provide a query or document IDs")
                    continue

                cmd, query = parts

                if cmd == "search":
                    search_result = await self.search_chain.run(
                        query, output_format="structured"
                    )
                    result = search_result["context"]
                    print(f"\nSearch Results:\n{result}")
                elif cmd == "summary":
                    search_result = await self.search_chain.run(
                        query, output_format="summary"
                    )
                    result = search_result["context"]
                    print(f"\nSummary:\n{result}")
                elif cmd == "analysis":
                    search_result = await self.search_chain.run(
                        query, output_format="analysis"
                    )
                    result = search_result["context"]
                    print(f"\nAnalysis:\n{result}")
                elif cmd == "get":
                    result = await self.search_chain.get_document_context_tool(query)
                    print(f"\nDocument Context:\n{result}")
                else:
                    print(
                        "Unknown command. Use 'search', 'summary', 'analysis', 'get', or 'quit'"
                    )

            except KeyboardInterrupt:
                break
            except Exception as e:
                print(f"Error: {e}")

        await self.search_chain.close()

    async def run_mcp_mode(self):
        """Run in MCP mode."""
        if not MCP_AVAILABLE:
            raise RuntimeError("MCP library not available")

        logger.info(f"Starting Azure AI Search MCP Server v{config.mcp_server.version}")

        try:
            # Use stdio transport
            from mcp.server.stdio import stdio_server

            async with stdio_server() as (read_stream, write_stream):
                from mcp.server import NotificationOptions
                from mcp.server.models import InitializationOptions

                initialization_options = InitializationOptions(
                    server_name=config.mcp_server.name,
                    server_version=config.mcp_server.version,
                    capabilities=self.server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={},
                    ),
                )

                await self.server.run(read_stream, write_stream, initialization_options)

        except asyncio.CancelledError:
            logger.info("Server run was cancelled")
            raise
        except Exception as e:
            import traceback
            print(f"[{time.strftime('%H:%M:%S')}] SERVER_TRACEBACK:", file=sys.stderr)
            traceback.print_exc(file=sys.stderr)
            logger.error(f"Server error: {str(e)}")
            raise
        finally:
            await self.search_chain.close()


def main():
    """Main entry point."""
    import argparse

    parser = argparse.ArgumentParser(description="Azure AI Search MCP Server")
    parser.add_argument(
        "--standalone", action="store_true", help="Run in standalone mode for testing"
    )
    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
    args = parser.parse_args()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    async def async_main():
        server = AzureSearchMCPServer()

        # Determine mode based on arguments
        if args.standalone or not MCP_AVAILABLE:
            # Standalone mode for testing
            logger.info("Running in standalone mode")
            await server.run_standalone_mode()
        else:
            # MCP mode (default)
            logger.info("Running in MCP mode (awaiting client connection)")
            await server.run_mcp_mode()

    try:
        asyncio.run(async_main())
    except KeyboardInterrupt:
        logger.info("Server stopped by user")
    except Exception as e:
        logger.error(f"Server failed: {e}")
        raise


if __name__ == "__main__":
    main()