"""
ProblemRetrieverAgent - ищет задачи в базе данных по теме.
"""

import logging
import json
from typing import List, Any
from google.genai import types
from .base_agent import BaseAgent, AgentState, AgentResponse
from mcpserver.mcp_client import get_llm_function_definitions

logger = logging.getLogger(__name__)

class ProblemRetrieverAgent(BaseAgent):
    def __init__(self):
        system_prompt = (
            "Твоя роль: AI-ассистент для поиска учебных задач. Ты должен строго следовать правилам.\n\n"
            "**ОБЯЗАТЕЛЬНЫЙ АЛГОРИТМ:**\n"
            "1. **ПОЛУЧИ СПИСОК ТЕМ:** Вызови `llm_list_available_topics`.\n"
            "2. **ВЫБЕРИ ЛУЧШУЮ ТЕМУ:** Сравни запрос пользователя с темами из списка и выбери одну, наиболее близкую по смыслу.\n"
            "3. **ЗАПРОСИ ЗАДАЧУ:** Вызови `llm_get_problems_by_topic` с `topic_id` выбранной темы. В результате ты получишь JSON с данными задачи (проигнорируй любые поля, содержащие ответ).\n"
            "4. **СФОРМИРУЙ ВЫВОД:** Твой единственный вывод в чат: только текст задачи.\n\n"
            "5. **ОТМЕТЬ ВЫДАННУЮ ЗАДАЧУ:** Dызови `llm_mark_problem_as_solved` с `problem_id` из полученного JSON, чтобы задача записалась как выданная конкретному пользователю.\n"
            "6. НАПИШИ ТОЛЬКО ТЕКСТ ЗАДАЧИ, НИКАКИХ ДРУГИХ СЛОВ, ПРИВЕТОВ, ПОЯСНЕНИЙ ИЛИ ДРУГИХ СЛОВ.\n"

            "**ПРАВИЛА:**\n"
            "- НИКОГДА НЕ ДОБАВЛЯЙ СЛОВА, ПРИВЕТОВ, ПОЯСНЕНИЙ ИЛИ ДРУГИХ СЛОВ, НИЧЕГО, КРОМЕ ТЕКСТА ЗАДАЧИ.\n"
            "- Не добавляй скрытые ответы или HTML-комментарии.\n"
            "- Не передавай и не упоминай ответ на задачу из БД — он запрещён к использованию.\n"
            "- Текст задачи должен использовать LaTeX-разметку для математических выражений.\n"
            "- Формулы оформляй строго с разделителями: для встроенных формул \\( ... \\), для блочных — \\[ ... \\]. Не используй $...$, $$...$$ и код-блоки ```latex```/```math``` для формул.\n"
            "- Если исходный текст содержит другие разделители (например, $...$, $$...$$), перепиши их в указанный формат \\( ... \\) или \\[ ... \\], сохраняя математический смысл. Не помещай формулы в кодовые блоки и не используй HTML."
        )
        super().__init__("ProblemRetriever", system_prompt, mode=0, history_limit=7)
        # Subjects where DB-driven problem retrieval is disabled
        self._disabled_subjects = {3}  # geometry

    def _get_tools(self) -> List[Any]:
        """Compatibility method - not used"""
        return []

    def _get_gemini_tools(self) -> List[Any]:
        """Get MCP tools in Gemini format"""
        try:
            # Disable tools entirely for disabled subjects (state guard in BaseAgent will short-circuit anyway)
            # Tools are static here; additional guard exists in BaseAgent._send_message_gemini
            # MCP client returns a list of function declaration dicts.
            # We need to wrap it in a Tool object for Gemini.
            function_declarations = get_llm_function_definitions()
            return [types.Tool(function_declarations=function_declarations)]
        except Exception as e:
            logger.error(f"ProblemRetriever: Error getting Gemini MCP tools: {e}")
            return []

    def _process_function_call(self, function_call, state: AgentState) -> Any:
        """Process MCP function calls"""
        try:
            # Block any tool usage for disabled subjects
            if int(getattr(state, 'subject_id', 0)) in getattr(self, '_disabled_subjects', set()):
                # Direct response to inform the user succinctly
                return {"__direct_response__": "Для геометрии выдача задач из базы отключена. Пришлите свою задачу или запросите объяснение темы."}

            from mcpserver.mcp_client import (
                llm_list_available_topics, 
                llm_get_problems_by_topic, 
                llm_mark_problem_as_solved
            )
            
            func_name = function_call.name
            
            args = {}
            if hasattr(function_call, 'args') and function_call.args is not None:
                args = dict(function_call.args)
            elif hasattr(function_call, 'input') and function_call.input is not None:
                args = function_call.input
            
            user_id = state.user_id
            if not user_id:
                logger.error("ProblemRetriever: user_id is not in agent state!")
                return "Ошибка: Не удалось определить пользователя. Пожалуйста, попробуйте еще раз."

            if func_name == "llm_get_problems_by_topic":
                # Always use subject_id from state to avoid spoofing from LLM args
                subject_id = int(state.subject_id)
                topic_id = args.get("topic_id")  # Changed from topic_name to topic_id
                
                if topic_id is None:
                    return "Ошибка: Не указан ID темы. Пожалуйста, предоставьте ID темы."
                
                problems_json_str = llm_get_problems_by_topic(subject_id, topic_id, user_id)
                problems_data = json.loads(problems_json_str)
                
                if problems_data.get("status") == "success" and problems_data.get("problems"):
                    # Take first unsolved problem
                    raw_problem = dict(problems_data["problems"][0])
                    problem_id = raw_problem.get("id")
                    # Prepare plain task text only
                    task_text = raw_problem.get("question", "")
                    if not isinstance(task_text, str):
                        task_text = json.dumps(task_text, ensure_ascii=False)

                    # Auto-mark as issued/solved for this user
                    try:
                        if problem_id is not None:
                            _ = llm_mark_problem_as_solved(user_id, int(problem_id))
                            logger.info(f"ProblemRetriever: Auto-marked problem {problem_id} as issued for user {user_id}.")
                    except Exception as e:
                        logger.error(f"ProblemRetriever: Error auto-marking problem {problem_id}: {e}")

                    # Return direct response with ONLY the task text
                    return {"__direct_response__": task_text}
                else:
                    return f"К сожалению, я не нашел задач по теме (ID: {topic_id}). Попробуйте другую тему."
            
            elif func_name == "llm_mark_problem_as_solved":
                problem_id = int(args.get("problem_id"))
                # Always use the user_id from the state for security and correctness
                
                mark_result_str = llm_mark_problem_as_solved(user_id, problem_id)
                mark_result = json.loads(mark_result_str)

                if mark_result.get("status") == "success":
                    logger.info(f"Problem {problem_id} marked as solved for user {user_id}.")
                    return f"Задача {problem_id} отмечена как выданная."
                else:
                    logger.error(f"Failed to mark problem {problem_id} as solved for user {user_id}. Reason: {mark_result.get('message')}")
                    return f"Не удалось отметить задачу {problem_id} как выданную. Причина: {mark_result.get('message')}"

            # Keep other function calls as they are               
            elif func_name == "llm_list_available_topics":
                # Always use subject_id from state
                subject_id = int(state.subject_id)
                return llm_list_available_topics(subject_id)
                
        except Exception as e:
            logger.error(f"ProblemRetriever: Error processing function call {getattr(function_call, 'name', 'unknown')}: {e}", exc_info=True)
            return "Произошла ошибка при поиске задач. Пожалуйста, попробуйте еще раз. Возможно, возникла проблема с идентификатором пользователя."
        
        return None