"""
Multi-agent orchestration module.
Coordinates RouterAgent and specialized agents in synchronous flow.
"""

import logging
from typing import Dict, Type

from message_utils import get_last_messages # Import get_last_messages

from agents import (
    RouterAgent,
    ExplainerAgent,
    SolverAgent,
    ProblemRetrieverAgent,
    ButtonSuggesterAgent,
    AgentState,
    AgentResponse,
)

logger = logging.getLogger(__name__)


class AgentRegistry:
    """Holds singletons of agent instances"""

    def __init__(self, api_key: str):
        self.api_key = api_key
        self._agents: Dict[str, object] = {}
        self._initialize_agents()

    def _initialize_agents(self):
        self._agents = {
            "RouterAgent": RouterAgent(),
            "ExplainerAgent": ExplainerAgent(),
            "SolverAgent": SolverAgent(),
            "ProblemRetrieverAgent": ProblemRetrieverAgent(),
            "ButtonSuggesterAgent": ButtonSuggesterAgent(),
        }
        logger.info(f"AgentRegistry initialized with {len(self._agents)} agents")

    def get(self, name: str):
        return self._agents.get(name)

    def list_agents(self):
        return list(self._agents.keys())


class Orchestrator:
    """Simple synchronous orchestrator for agent interaction"""

    MAX_HOPS = 5  # Prevent infinite loops

    def __init__(self, api_key: str = None): # Removed GEMINI_API_KEY
        self.registry = AgentRegistry(api_key)

    def handle_message(self, message: str, state: AgentState) -> AgentResponse:
        """Process user message through Router and specialized agents"""
        hop = 0
        current_agent_name = state.context.get("current_agent", "RouterAgent")
        last_response: AgentResponse = None

        while hop < self.MAX_HOPS and current_agent_name:
            hop += 1
            agent = self.registry.get(current_agent_name)
            if agent is None:
                logger.error(f"Unknown agent: {current_agent_name}")
                break

            logger.info(f"Orchestrator: Hop {hop} -> {current_agent_name}")

            # Fetch chat history based on agent's limit and pass it explicitly
            history_limit = getattr(agent, 'history_limit', 100)
            # Fetch up to the agent's history limit from DB
            history_to_fetch = history_limit if history_limit > 0 else 0
            
            history = []
            # Prefer externally supplied history (e.g., from Telegram in-memory chat)
            external_history = None
            try:
                external_history = state.context.get("external_history")
            except Exception:
                external_history = None

            if isinstance(external_history, list) and history_to_fetch > 0:
                # Keep only the most recent items within the limit minus the current user message
                slice_len = max(0, history_to_fetch - 1)
                history = external_history[-slice_len:] if slice_len > 0 else []
            elif state.chat_id and history_to_fetch > 0:
                history = get_last_messages(state.chat_id, limit=history_to_fetch)

            # Avoid duplicating the current user message if it's already saved in DB
            # The 'message' can be a JSON string for multimodal content
            if not history or history[-1].get("role") != "user" or history[-1].get("content") != message:
                history.append({"role": "user", "content": message})
            
            # Call agent with the full history
            response = agent.process_message(history, state)

            # Update state
            if response.update_state:
                for k, v in response.update_state.items():
                    setattr(state, k, v)

            # If there's an actual content from the agent, save it.
            # We no longer append to state.chat_history directly as it's not part of AgentState
            # state.chat_history.append({"role": "assistant", "content": response.content})

            # Decide next agent
            if response.next_agent:
                current_agent_name = response.next_agent
            else:
                last_response = response
                break

        # After conversation step, call LoggerAgent side effect
        try:
            logger_agent = self.registry.get("LoggerAgent")
            if logger_agent:
                logger_agent.process_message(message, state)
        except Exception as e:
            logger.error(f"LoggerAgent error: {e}")

        if last_response is None:
            last_response = AgentResponse(agent_name="system", content="Ошибка маршрутизации", is_final=True)
        return last_response


# Helper singleton orchestrator
_global_orchestrator: Orchestrator = None

def get_orchestrator() -> Orchestrator:
    global _global_orchestrator
    if _global_orchestrator is None:
        _global_orchestrator = Orchestrator() # API key is now handled by individual agents
    return _global_orchestrator