"""Agent orchestrator for managing AI conversations.

This module provides the main orchestrator that ties together
models, channels, tools, and memory into a complete agent.
"""

import logging
from typing import Any, Dict, List, Optional

from agent_factory.channels.base import BaseChannel, ChannelResponse
from agent_factory.core.gateway import ToolGateway
from agent_factory.core.memory import Memory, MemoryStore
from agent_factory.models.base import (
    BaseModel,
    Message,
    ModelResponse,
    ToolDefinition,
)
from agent_factory.tools.base import BaseTool

logger = logging.getLogger(__name__)


class AgentOrchestrator:
    """Main orchestrator for AI agents.

    Coordinates between AI models, communication channels, tools,
    and conversation memory to handle complete agent interactions.

    Example:
        from agent_factory import AgentOrchestrator
        from agent_factory.models import BedrockModel
        from agent_factory.channels import SlackChannel

        # Create agent
        agent = AgentOrchestrator(
            model=BedrockModel(),
            system_prompt="You are a helpful assistant.",
        )

        # Register tools
        agent.register_tool(my_tool)

        # Handle a message
        response = agent.invoke("Hello!")

        # Or handle channel messages
        agent.add_channel(SlackChannel())
        agent.handle_channel_event(slack_event)
    """

    def __init__(
        self,
        model: BaseModel,
        system_prompt: str = "You are a helpful AI assistant.",
        tools: Optional[List[BaseTool]] = None,
        channels: Optional[List[BaseChannel]] = None,
        memory_store: Optional[MemoryStore] = None,
        fallback_model: Optional[BaseModel] = None,
        max_tool_iterations: int = 10,
    ):
        """Initialize the agent orchestrator.

        Args:
            model: Primary AI model to use
            system_prompt: System prompt for the agent
            tools: Initial tools to register
            channels: Communication channels
            memory_store: Storage backend for memory
            fallback_model: Fallback model if primary fails
            max_tool_iterations: Maximum tool use iterations
        """
        self.model = model
        self.fallback_model = fallback_model
        self.system_prompt = system_prompt
        self.max_tool_iterations = max_tool_iterations

        # Initialize components
        self.gateway = ToolGateway()
        self._memory_store = memory_store
        self._memory_cache: Dict[str, Memory] = {}
        self._channels: Dict[str, BaseChannel] = {}

        # Register initial tools
        if tools:
            for tool in tools:
                self.gateway.register(tool)

        # Register channels
        if channels:
            for channel in channels:
                self.add_channel(channel)

    def get_memory(self, session_id: str) -> Memory:
        """Get or create memory for a session.

        Args:
            session_id: Session identifier

        Returns:
            Memory instance for the session
        """
        if session_id not in self._memory_cache:
            self._memory_cache[session_id] = Memory(
                session_id=session_id,
                store=self._memory_store,
            )
        return self._memory_cache[session_id]

    def register_tool(self, tool: BaseTool) -> None:
        """Register a tool with the agent.

        Args:
            tool: Tool to register
        """
        self.gateway.register(tool)

    def add_channel(self, channel: BaseChannel) -> None:
        """Add a communication channel.

        Args:
            channel: Channel to add
        """
        self._channels[channel.channel_type.value] = channel

    def invoke(
        self,
        prompt: str,
        session_id: str = "default",
        context: Optional[Dict[str, Any]] = None,
        include_history: bool = True,
    ) -> Dict[str, Any]:
        """Invoke the agent with a prompt.

        Args:
            prompt: User prompt
            session_id: Session ID for memory
            context: Context to inject into tools
            include_history: Whether to include conversation history

        Returns:
            Response dictionary with result and metadata
        """
        memory = self.get_memory(session_id)
        context = context or {}

        # Build messages
        messages: List[Message] = []
        if include_history:
            messages = memory.get_messages()

        # Add user message
        messages.append(Message(role="user", content=prompt))

        # Get tools
        tools = self.gateway.get_tool_definitions()

        # Converse with tool loop
        response = self._converse_with_tools(
            messages=messages,
            tools=tools,
            context=context,
        )

        # Save to memory
        memory.add_turn(prompt, response.content)

        return {
            "result": response.content,
            "session_id": session_id,
            "model_id": response.model_id,
            "usage": response.usage,
        }

    def _converse_with_tools(
        self,
        messages: List[Message],
        tools: List[ToolDefinition],
        context: Optional[Dict[str, Any]] = None,
    ) -> ModelResponse:
        """Converse with AI including tool execution loop.

        Args:
            messages: Conversation messages
            tools: Available tools
            context: Context to inject into tools

        Returns:
            Final model response
        """
        context = context or {}
        iterations = 0

        # Get initial response
        response = self._call_model(messages, tools)

        # Tool use loop
        while response.has_tool_use and iterations < self.max_tool_iterations:
            iterations += 1
            tool_results = []

            for tool_use in response.tool_uses:
                logger.info(f"Executing tool: {tool_use.name}")

                # Execute tool with context injection
                result = self.gateway.execute(tool_use, context=context)
                tool_results.append(result)

                # Add tool use to messages
                messages.append(
                    Message(
                        role="assistant",
                        content="",
                        tool_use=tool_use,
                    )
                )

                # Add tool result to messages
                messages.append(
                    Message(
                        role="user",
                        content="",
                        tool_result=result.to_bedrock_format(),
                    )
                )

            # Get next response
            response = self._call_model(messages, tools)

        return response

    def _call_model(
        self,
        messages: List[Message],
        tools: List[ToolDefinition],
    ) -> ModelResponse:
        """Call the AI model with fallback.

        Args:
            messages: Conversation messages
            tools: Available tools

        Returns:
            Model response

        Raises:
            ModelError: If all models fail
        """
        from agent_factory.models.base import ModelError

        # Try primary model
        try:
            return self.model.converse(
                messages=messages,
                system_prompt=self.system_prompt,
                tools=tools if tools else None,
            )
        except ModelError as e:
            if not e.is_transient or not self.fallback_model:
                raise

            logger.warning(f"Primary model failed, trying fallback: {e}")

        # Try fallback model
        if self.fallback_model and self.fallback_model.is_available():
            return self.fallback_model.converse(
                messages=messages,
                system_prompt=self.system_prompt,
                tools=tools if tools else None,
            )

        raise ModelError("All models failed", model_id=self.model.model_id)

    def handle_channel_event(
        self,
        event: Dict[str, Any],
        channel_type: Optional[str] = None,
    ) -> Optional[ChannelResponse]:
        """Handle an event from a channel.

        Args:
            event: Raw event from the channel
            channel_type: Channel type (auto-detected if not provided)

        Returns:
            Channel response or None
        """
        # Find channel
        channel = None
        if channel_type:
            channel = self._channels.get(channel_type)
        else:
            # Try to detect channel from event
            for ch in self._channels.values():
                try:
                    if ch.parse_event(event):
                        channel = ch
                        break
                except Exception:
                    continue

        if not channel:
            logger.warning("No channel found for event")
            return None

        # Parse message
        message = channel.parse_event(event)
        if not message:
            return None

        # Get context from channel
        context = channel.get_context(message)

        # Process message
        result = self.invoke(
            prompt=message.text,
            session_id=message.session_id,
            context=context,
        )

        # Format response
        response = channel.format_response(
            text=result["result"],
            channel_id=message.channel_id,
            thread_id=message.thread_id,
        )

        # Send response
        channel.send_message(response)

        return response

    def clear_session(self, session_id: str) -> None:
        """Clear memory for a session.

        Args:
            session_id: Session to clear
        """
        if session_id in self._memory_cache:
            self._memory_cache[session_id].clear()
            del self._memory_cache[session_id]

    @property
    def tools(self) -> List[str]:
        """Get list of registered tool names."""
        return self.gateway.list_tools()
