import asyncio
import yaml
from functools import cached_property

from kwark.command import KwarkCommand
from kwark.ai_services import AIService
from kwark.ai_services.anthropic_ai_service import (
    AnthropicAIService, AnthropicToolset
)
from kwark.ai_services.mcp_client import MCPClientWrapper


class ChatCommand(KwarkCommand):
    """Simple chat command to interact with AI.
    Takes input text and returns an AI response."""

    name = 'chat'

    def _parse_mcp_config(self, config_text):
        """Parse MCP server configuration from YAML text.

        Args:
            config_text: YAML string containing MCP server definitions

        Returns:
            List of server dicts with 'command' and 'args' keys
        """
        if not config_text or not config_text.strip():
            return []

        config = yaml.safe_load(config_text)
        if not config or 'mcp' not in config:
            return []

        servers = []
        for server_line in config['mcp']:
            parts = server_line.split()
            if parts:
                servers.append({
                    'command': parts[0],
                    'args': parts[1:] if len(parts) > 1 else []
                })
        return servers

    def _init_mcp_clients(self, server_configs):
        """Initialize MCP clients from server configurations

        Args:
            server_configs: List of server config dicts

        Returns:
            List of connected MCPClientWrapper instances
        """
        if not server_configs:
            return []

        clients = []

        async def connect_all():
            """Connect to all MCP servers"""
            tasks = []
            for cfg in server_configs:
                wrapper = MCPClientWrapper(cfg['command'], cfg['args'])
                tasks.append(self._connect_and_discover(wrapper))
            return await asyncio.gather(*tasks, return_exceptions=True)

        # Use the persistent event loop
        results = self.mcp_loop.run_until_complete(connect_all())

        for i, result in enumerate(results):
            if isinstance(result, Exception):
                self.app.ui.send(
                    f"Warning: Failed to connect to MCP server: {result}",
                    emphasis='error')
            elif result:
                clients.append(result)

        return clients

    async def _connect_and_discover(self, wrapper):
        """Connect to MCP server and discover tools

        Args:
            wrapper: MCPClientWrapper instance

        Returns:
            wrapper with tools populated, or None on error
        """
        try:
            await wrapper.connect()
            wrapper.tools = await wrapper.list_tools()
            return wrapper
        except Exception as e:
            raise Exception(
                f"{wrapper.command} {' '.join(wrapper.args)}: {str(e)}")

    def handle_vals(self):
        """Override to parse MCP config and set up toolset"""
        super().handle_vals()

        # Parse MCP config from stream (before it's used elsewhere)
        mcp_config_text = self.app.stream.text
        mcp_config = self._parse_mcp_config(mcp_config_text)

        # Only initialize MCP if config present
        if mcp_config:
            # Create event loop for MCP operations
            self.mcp_loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self.mcp_loop)

            # Initialize MCP clients
            mcp_clients = self._init_mcp_clients(mcp_config)

            # Create toolset with MCP clients
            try:
                self.toolset = AnthropicToolset(mcp_clients=mcp_clients)
            except ValueError as e:
                self.app.ui.send(f"Error: {e}", emphasis='error')
                self.toolset = AnthropicToolset()
        else:
            self.toolset = None  # Use default
            self.mcp_loop = None

    @cached_property
    def ai_service(self):
        """Override to inject MCP toolset if present"""
        # Use parent's ai_service creation (respects mocking)
        service = AIService.create(self.ai_service_type, self.api_key)

        # Inject MCP toolset and event loop if we created them
        if hasattr(self, 'toolset') and self.toolset:
            service.toolset = self.toolset
            service.mcp_loop = self.mcp_loop

        return service

    @KwarkCommand.wrap
    def execute(self):
        # If we have MCP config, stream was consumed for that
        # Otherwise, stream text can be initial message
        if hasattr(self, 'toolset') and self.toolset:
            initial_message = None
        else:
            initial_message = self.app.stream.text.strip() or None

        try:
            self.ai_service.chat(self.app.ui, initial_message=initial_message)
        finally:
            # Clean up MCP clients
            if hasattr(self, 'toolset') and self.toolset:
                self._cleanup_mcp_clients()

        self.status = "Chat session completed"

    def _cleanup_mcp_clients(self):
        """Clean up MCP client connections"""
        if not self.toolset or not self.toolset.mcp_clients:
            return

        async def cleanup():
            for client in self.toolset.mcp_clients:
                try:
                    await client.cleanup()
                except Exception:
                    pass  # Ignore cleanup errors

        try:
            if self.mcp_loop and not self.mcp_loop.is_closed():
                self.mcp_loop.run_until_complete(cleanup())
                self.mcp_loop.close()
        except Exception:
            pass  # Ignore cleanup errors
