"""MCP prompt commands for ACP slash command integration."""

from __future__ import annotations

from typing import TYPE_CHECKING

from acp.schema import AvailableCommand, AvailableCommandInput, CommandInputHint
from llmling_agent.log import get_logger
from llmling_agent_acp.converters import to_agent_text_notification


if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from mcp.types import Prompt as MCPPrompt

    from acp.schema import SessionNotification
    from llmling_agent_acp.session import ACPSession


logger = get_logger(__name__)


class MCPPromptCommand:
    """Wrapper for MCP prompts as slash commands."""

    def __init__(self, mcp_prompt: MCPPrompt) -> None:
        """Initialize with MCP prompt.

        Args:
            mcp_prompt: MCP prompt object from server
        """
        self.mcp_prompt = mcp_prompt
        self.name = mcp_prompt.name
        self.description = mcp_prompt.description or f"MCP prompt: {mcp_prompt.name}"

    def to_available_command(self) -> AvailableCommand:
        """Convert to ACP AvailableCommand format.

        Returns:
            ACP AvailableCommand object
        """
        # Create input spec from MCP prompt arguments
        spec = None
        if self.mcp_prompt.arguments:
            arg_names = [arg.name for arg in self.mcp_prompt.arguments]
            hint = f"Arguments: {', '.join(arg_names)}"
            spec = AvailableCommandInput(root=CommandInputHint(hint=hint))
        name = f"mcp-{self.name}"  # Prefix to avoid conflicts
        return AvailableCommand(name=name, description=self.description, input=spec)

    async def execute(
        self,
        args: str,
        session: ACPSession,
    ) -> AsyncIterator[SessionNotification]:
        """Execute MCP prompt command.

        Args:
            args: Command arguments string
            session: ACP session context

        Yields:
            SessionNotification objects with prompt results
        """
        try:
            # Parse arguments if needed
            arguments = self._parse_arguments(args) if args.strip() else None

            # Get MCP manager from session
            if not session.mcp_manager:
                error_msg = "No MCP servers available"
                if update := to_agent_text_notification(error_msg, session.session_id):
                    yield update
                return

            # Find appropriate MCP client (use first available for now)
            if not session.mcp_manager.clients:
                error_msg = "No MCP clients connected"
                if update := to_agent_text_notification(error_msg, session.session_id):
                    yield update
                return

            # Execute prompt via first available MCP client
            client = next(iter(session.mcp_manager.clients.values()))

            try:
                # Try with arguments first, fallback to no arguments
                try:
                    result = await client.get_prompt(self.mcp_prompt.name, arguments)
                except Exception as e:
                    if arguments:
                        logger.warning(
                            "MCP prompt with arguments failed, trying without: %s", e
                        )
                        result = await client.get_prompt(self.mcp_prompt.name)
                    else:
                        raise

                # Convert prompt result to text
                content_parts = []
                for message in result.messages:
                    if hasattr(message.content, "text"):
                        content_parts.append(message.content.text)  # type: ignore
                    else:
                        content_parts.append(str(message.content))

                output = "\n".join(content_parts)

                # Add argument info if provided
                if arguments:
                    arg_info = ", ".join(f"{k}={v}" for k, v in arguments.items())
                    output = (
                        f"Prompt '{self.mcp_prompt.name}' with "
                        f"args ({arg_info}):\n\n{output}"
                    )

                # Stream as session updates
                if update := to_agent_text_notification(output, session.session_id):
                    yield update

            except Exception as e:
                error_msg = f"MCP prompt execution failed: {e}"
                logger.exception("MCP prompt execution error")
                if update := to_agent_text_notification(error_msg, session.session_id):
                    yield update

        except Exception as e:
            error_msg = f"Command error: {e}"
            logger.exception("MCP command execution error")
            if update := to_agent_text_notification(error_msg, session.session_id):
                yield update

    def _parse_arguments(self, args_str: str) -> dict[str, str]:
        """Parse argument string to dictionary.

        Args:
            args_str: Raw argument string

        Returns:
            Dictionary of argument name to value
        """
        # Simple parsing - split on spaces and match to prompt arguments
        if not self.mcp_prompt.arguments:
            return {}

        args_list = args_str.strip().split()
        arguments = {}

        # Map positional arguments to prompt argument names
        for i, arg_value in enumerate(args_list):
            if i < len(self.mcp_prompt.arguments):
                arg_name = self.mcp_prompt.arguments[i].name
                arguments[arg_name] = arg_value

        return arguments
