"""
Router Agent - First agent that receives user messages and routes them to appropriate specialized agents
"""

import google.generativeai as genai
from google.generativeai.types import Tool
from typing import List, Dict, Any
import logging
import json
import time # Import time for logging
from .base_agent import BaseAgent, AgentState, AgentResponse

logger = logging.getLogger(__name__)


class RouterAgent(BaseAgent):
    """
    Router Agent analyzes user messages and determines which specialized agent should handle the request
    
    Routes to:
    - Explainer: Theory explanations, concept clarifications
    - Solver: Mathematical problems, physics calculations  
    - ProblemRetriever: Finding specific problems/exercises
    """
    
    def __init__(self):
        system_prompt = """
        Ты - Router Agent в системе AI-репетитора для школьников. Твоя задача - анализировать сообщения учеников и определять, какой специализированный агент должен обработать запрос.

        Текущий предмет (subject_id): {subject_id}
        
        Доступные агенты:
        1. EXPLAINER - объясняет теорию, концепции, определения, решенные задачи (которые уже решены)
        2. SOLVER - решает задачи И ДАЕТ ПОДСКАЗИК ПО РЕШЕНИЮ ДЛЯ УЧЕНИКОВ
        3. PROBLEM_RETRIEVER - находит подходящие задачи и упражнения по теме, чтобы ученик научился решать задачи

        Критерии для выбора агента:

        PROBLEM_RETRIEVER — ВЫБИРАЙ ЭТОГО АГЕНТА ТОЛЬКО В СЛЕДУЮЩИХ СЛУЧАЯХ:

        ПОЛЬЗОВАТЕЛЬ ЯВНО ИЛИ НЕЯВНО ПРОСИТ ПРЕДОСТАВИТЬ ЗАДАЧИ ДЛЯ РЕШЕНИЯ:
        - ПОСЛЕ ОБЪЯСНЕНИЯ ТЕМЫ УЧЕНИКУ БЫЛО ПРЕДЛОЖЕНО ПОТРЕНИРОВАТЬСЯ ИЛИ НАУЧИТСЯ РЕШАТЬ ЗАДАЧИ ПО ЭТОЙ ТЕМЕ, А УЧЕНИК СОГЛАСИЛСЯ
        - Формулирует прямой запрос: "Дай задачу", "Хочу решать", "Можно пример задачи?"
        - Просит задачи по теме с целью тренировки, практики или самопроверки
        - Запрашивает дополнительные задачи: "Дай ещё", "Хочу порешать больше"
        - Выражает готовность решать задачи после объяснения темы: "Давай попробую", "Теперь хочу потренироваться"
        - Хочет проверить знания по теме через задачи: "Насколько я понял?", "Можешь проверить меня?"
        - Хочет научиться решать задачи на практике
        - Соглашается на предложение решить задачу, даже если сам не инициировал
        

        EXPLAINER - если ученик:
        - Просит объяснить теорию, концепцию, определение, РЕШЕННУЮ задачу
        - Спрашивает "что такое...", "как работает...", "объясни...", "как ты решил эту задачу
        - Хочет понять теоретические основы

        SOLVER - если ученик:
        - просто прислал условие задачи
        - попросил решить задачу
        - Присылает задачу
        - Просит помочь с вычислениями
        - решает задчу сам, но просит дать подсказку или помочь или не знает как решить
        - Говорит "реши задачу", "помоги решить", "как найти..."

        
        ОБЯЗАТЕЛЬНЫЕ ПРАВИЛА:
        0. ЕСЛИ subject_id == 3, 12, 4, 13 (геометрия или информатика или химия или статистика) — НИКОГДА не выбирай PROBLEM_RETRIEVER.
        1. **СУПЕР-ВАЖНОЕ ПРАВИЛО:** Если в "Предыдущем ответе ассистента" содержится "problem_retriever", а ученик просит решить задачу, показать ответ или говорит "не знаю, как решать" - **ВСЕГДА** выбирай **SOLVER**. Это правило имеет наивысший приоритет.
        2. ВСЕГДА используй функцию route_to_agent для каждого сообщения ученика
        3. НЕ отвечай текстом - ТОЛЬКО вызывай функцию route_to_agent
        4. Даже если не уверен - выбери наиболее подходящий агент и вызови функцию
        5. ЗАПРЕЩЕНО отвечать обычным текстом без вызова функции
        
        ПРИМЕРЫ:
        Ученик: "Реши задачу: U=12В, R=4 Ом, найди I"
        → Вызови route_to_agent с target_agent="SOLVER"
        
        Ученик: "Что такое электрический ток?"
        → Вызови route_to_agent с target_agent="EXPLAINER"
        
        ПОМНИ: КАЖДОЕ сообщение ученика = ОДИН вызов route_to_agent!
        """
        
        super().__init__("Router", system_prompt, mode=0, history_limit=3)
    
    def _get_tools(self) -> List:
        """
        Get tools available to this agent.
        Since RouterAgent only uses Gemini tools for routing, we simply
        delegate to _get_gemini_tools.
        """
        return self._get_gemini_tools()

    def _get_gemini_tools(self) -> List:
        """Get routing tools in Gemini format"""
        # Using dicts for tool definition to avoid proto/enum issues
        # that cause pydantic validation errors.
        routing_tool_dict = {
            "function_declarations": [
                {
                    "name": "route_to_agent",
                    "description": "Route user message to appropriate specialized agent",
                    "parameters": {
                        "type": "OBJECT",
                        "properties": {
                            "target_agent": {
                                "type": "STRING",
                                "description": "Target agent name: EXPLAINER, SOLVER, or PROBLEM_RETRIEVER"
                            }
                        },
                        "required": ["target_agent"]
                    }
                }
            ]
        }
        
        return [routing_tool_dict]


    
    def _process_function_call(self, function_call, state: AgentState) -> Any:
        """Process routing function call"""
        if function_call.name == "route_to_agent":
            # Handle Gemini function call format
            args = dict(function_call.args)
                
            target_agent = args.get("target_agent", "").upper()
            reasoning = args.get("reasoning", "")
            topic = args.get("extracted_topic", "")
            
            current_timestamp = time.time()
            logger.info(f"Router: ⚙️ START PROCESSING ROUTING DECISION at {current_timestamp:.2f}:") # Updated log
            logger.info(f"Router: → Target Agent: {target_agent}")

            if topic:
                logger.info(f"Router: → Extracted Topic: {topic}")
            
            # Map agent names to actual agent classes
            agent_mapping = {
                "EXPLAINER": "ExplainerAgent",
                "SOLVER": "SolverAgent", 
                "PROBLEM_RETRIEVER": "ProblemRetrieverAgent"
            }

            # Hard block: geometry subject (3) must not route to ProblemRetriever
            if int(getattr(state, 'subject_id', 0)) == 3 and target_agent == "PROBLEM_RETRIEVER":
                logger.info("Router: subject_id=3 -> overriding PROBLEM_RETRIEVER to EXPLAINER")
                target_agent = "EXPLAINER"
            
            if target_agent not in agent_mapping:
                logger.warning(f"Router: ⚠️  Unknown target agent '{target_agent}', defaulting to EXPLAINER")
                target_agent = "EXPLAINER"
                logger.info(f"Router: → Final Target Agent: EXPLAINER (fallback)")
            
            final_agent = agent_mapping[target_agent]
            logger.info(f"Router: 🗺️ AGENT MAPPED TO: {final_agent} at {current_timestamp:.2f}") # Added log
            current_timestamp = time.time()
            logger.info(f"Router: 🎯 ROUTING COMPLETE at {current_timestamp:.2f}")
            
            return {
                "next_agent": final_agent
            }
        
        return None
    
    
    def process_message(
        self,
        history: List[Dict[str, Any]],
        state: AgentState
    ) -> AgentResponse:
        """Route message to appropriate agent"""
        current_timestamp = time.time()
        logger.info(f"Router: 🏁 START ROUTING ANALYSIS at {current_timestamp:.2f}:") # Updated log
        logger.info(f"Router: → Subject ID: {state.subject_id}")
        logger.info(f"Router: → User ID: {state.user_id}")
        logger.info(f"Router: → Provider: {getattr(self, 'provider', 'unknown')}")

        
        # Динамически обновляем системный промпт с subject_id и последним сообщением
        dynamic_system_prompt = self.system_prompt.format(
            subject_id=state.subject_id
        )
        
        # Log full history that Router receives and relies on for routing
        # This helps to see exactly what messages influenced the routing decision
        try:
            logger.info("Router: 📜 HISTORY MESSAGES (oldest first):")
            if not history:
                logger.info("Router: (history is empty)")
            else:
                for idx, msg in enumerate(history):
                    role = (
                        (msg.get("role") if isinstance(msg, dict) else None)
                        or (msg.get("author") if isinstance(msg, dict) else None)
                        or (msg.get("speaker") if isinstance(msg, dict) else None)
                        or "unknown"
                    )
                    content = msg.get("content") if isinstance(msg, dict) else str(msg)

                    # Normalize different content shapes to plain text for logging
                    text = ""
                    try:
                        if isinstance(content, str):
                            text = content
                        elif isinstance(content, list):
                            pieces = []
                            for part in content:
                                if isinstance(part, dict):
                                    if "text" in part and isinstance(part["text"], str):
                                        pieces.append(part["text"])
                                    elif "content" in part and isinstance(part["content"], str):
                                        pieces.append(part["content"])
                                    else:
                                        pieces.append(json.dumps(part, ensure_ascii=False))
                                else:
                                    pieces.append(str(part))
                            text = "\n".join(pieces)
                        elif isinstance(content, dict):
                            if "text" in content and isinstance(content["text"], str):
                                text = content["text"]
                            elif "parts" in content and isinstance(content["parts"], list):
                                parts_text = []
                                for part in content["parts"]:
                                    if isinstance(part, dict) and "text" in part and isinstance(part["text"], str):
                                        parts_text.append(part["text"])
                                    else:
                                        parts_text.append(str(part))
                                text = "\n".join(parts_text)
                            else:
                                text = json.dumps(content, ensure_ascii=False)
                        else:
                            text = str(content)
                    except Exception as inner_e:
                        text = f"<unprintable content: {inner_e}>"

                    if text is None:
                        text = ""
                    text_to_log = text if len(text) <= 2000 else text[:2000] + "… [truncated]"
                    logger.info(f"Router: #{idx+1} {role}: {text_to_log}")
        except Exception as e:
            logger.error(f"Router: Error while logging history: {e}")

        # Временно обновляем системный промпт
        original_system_prompt = self.system_prompt
        self.system_prompt = dynamic_system_prompt
        
        response = super().process_message(history, state)
        logger.info(f"Router: 🧠 LLM DECISION PROCESS COMPLETED at {time.time():.2f}") # Added log
        
        # Restore original system prompt
        self.system_prompt = original_system_prompt
        
        # If no routing happened, default to Explainer
        if not response.next_agent:
            logger.warning("Router: ⚠️  No routing decision made by LLM")
            response.next_agent = "ExplainerAgent"
            
        
        # Router doesn't provide final content, just routes
        response.content = ""
        response.stream = None
        response.is_final = False
        
        logger.info(f"Router: ✅ END ROUTING ANALYSIS for message at {time.time():.2f}") # Added log
        return response 
