"""
Claude Agent SDK Integration - MCP Server and Options Export.

Provides utilities to export AMI tools in Claude-compatible format.
"""
import base64
import json
import logging
from typing import Any, Optional

from memrail.tools.registry import ToolRegistry


logger = logging.getLogger(__name__)


class ClaudeBridge:
    """
    Bridge between AMI Tool Registry and Claude Agent SDK.

    Exports registered tools as MCP server configuration and generates
    ClaudeAgentOptions-compatible configuration for seamless integration.

    Example:
        ```python
        bridge = ClaudeBridge(registry)

        # Get MCP server configuration
        mcp_server = bridge.to_mcp_server(name="ami_tools")

        # Get full Claude options
        claude_opts = bridge.to_claude_options(server_alias="ami")

        # Use with ClaudeSDKClient
        options = ClaudeAgentOptions(
            **claude_opts,
            permission_mode="acceptEdits",
        )
        ```
    """

    def __init__(self, registry: ToolRegistry):
        """
        Initialize the Claude bridge.

        Args:
            registry: ToolRegistry with registered tools
        """
        self.registry = registry

    def to_mcp_server(
        self,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> dict[str, Any]:
        """
        Export tools as an MCP server configuration.

        Returns a configuration that can be used with Claude Agent SDK's
        mcp_servers option.

        Args:
            name: Server name (defaults to registry server_name)
            description: Server description

        Returns:
            MCP server configuration dict
        """
        server_name = name or self.registry.server_name or "ami_tools"

        # Build tools configuration
        tools_config = self._build_tools_config()

        # Serialize to base64 for passing through environment
        config_json = json.dumps({
            "name": server_name,
            "description": description or f"AMI tools from {server_name}",
            "tools": tools_config,
        })
        config_b64 = base64.b64encode(config_json.encode()).decode()

        return {
            "type": "stdio",
            "command": "python",
            "args": ["-m", "memrail.tools.mcp_server"],
            "env": {
                "AMI_TOOLS_CONFIG": config_b64,
            },
        }

    def to_claude_options(
        self,
        server_alias: str = "ami",
        include_mcp_server: bool = True,
    ) -> dict[str, Any]:
        """
        Export full ClaudeAgentOptions-compatible configuration.

        Args:
            server_alias: Alias for the MCP server (tools will be named mcp__{alias}__toolname)
            include_mcp_server: Whether to include the mcp_servers configuration

        Returns:
            Dict with mcp_servers and allowed_tools configurations
        """
        result = {}

        if include_mcp_server:
            result["mcp_servers"] = {
                server_alias: self.to_mcp_server(),
            }

        # Generate allowed_tools list
        allowed_tools = []
        for tool in self.registry.list_tools():
            # MCP tool naming convention: mcp__{server}__{tool_name}
            allowed_tools.append(f"mcp__{server_alias}__{tool.name}")

        result["allowed_tools"] = allowed_tools

        return result

    def _build_tools_config(self) -> list[dict[str, Any]]:
        """Build the tools configuration for MCP."""
        tools = []

        for tool_def in self.registry.list_tools():
            tool_config = {
                "name": tool_def.name,
                "description": tool_def.description,
                "inputSchema": tool_def.to_json_schema(),
            }

            # Add AMI metadata as annotations
            # Always include projects in annotations so they stay in sync
            if tool_def.projects:
                tool_config["annotations"] = tool_config.get("annotations", {})
                tool_config["annotations"]["projects"] = tool_def.projects

            if tool_def.capabilities:
                tool_config["annotations"] = tool_config.get("annotations", {})
                tool_config["annotations"]["capabilities"] = tool_def.capabilities

            if tool_def.constraints:
                tool_config["annotations"] = tool_config.get("annotations", {})
                tool_config["annotations"]["constraints"] = tool_def.constraints

            tools.append(tool_config)

        return tools

    def generate_tool_use_prompt(self) -> str:
        """
        Generate a prompt section describing available tools.

        Useful for including in system prompts when AMI tools should be
        available to Claude.

        Returns:
            Formatted string describing available tools
        """
        lines = [
            "# Available AMI Tools",
            "",
            "The following tools are registered with AMI and can be executed:",
            "",
        ]

        for tool_def in self.registry.list_tools():
            lines.append(f"## {tool_def.name}")
            lines.append(f"**Description:** {tool_def.description}")
            lines.append("")

            if tool_def.capabilities:
                lines.append(f"**Capabilities:** {', '.join(tool_def.capabilities)}")
            if tool_def.constraints:
                lines.append("**Constraints:**")
                for constraint in tool_def.constraints:
                    lines.append(f"  - {constraint}")
            lines.append("")

            # Show input schema
            lines.append("**Input Schema:**")
            lines.append("```json")
            lines.append(json.dumps(tool_def.to_json_schema(), indent=2))
            lines.append("```")
            lines.append("")

        return "\n".join(lines)


def create_mcp_tool_handler(registry: ToolRegistry):
    """
    Create an MCP tool handler function for use with Claude Agent SDK.

    This handler can be used directly with Claude's tool_use callback
    to execute AMI-registered tools.

    Args:
        registry: ToolRegistry with registered tools

    Returns:
        Async function that handles tool calls

    Example:
        ```python
        handler = create_mcp_tool_handler(registry)

        # In Claude tool_use callback:
        result = await handler(tool_name, tool_args)
        ```
    """
    async def handle_tool_call(tool_name: str, args: dict) -> dict[str, Any]:
        """Handle a tool call from Claude."""
        tool = registry.get_tool(tool_name)

        if tool is None:
            return {
                "error": f"Tool not found: {tool_name}",
                "success": False,
            }

        if not tool.has_handler:
            return {
                "error": f"Tool '{tool_name}' has no handler bound",
                "success": False,
            }

        try:
            result = await tool.handler(args)
            return {
                "success": True,
                "result": result,
            }
        except Exception as e:
            logger.exception(f"Tool execution error: {tool_name}")
            return {
                "error": str(e),
                "success": False,
            }

    return handle_tool_call
