"""
Простой агент для Agent Challenge с использованием Gradio и Mixtral
"""

import os
import re
import math
import json
import gradio as gr
from huggingface_hub import InferenceClient

# Безопасная обработка токена Hugging Face
# Токен должен быть установлен как переменная окружения HF_TOKEN
# или передан через Secrets в Hugging Face Spaces
HF_TOKEN = os.environ.get("HF_TOKEN")

# Инициализация клиента Hugging Face
client = None
try:
    client = InferenceClient(
        model="mistralai/Mixtral-8x7B-Instruct-v0.1",  # Рекомендуемая модель
        token=HF_TOKEN
    )
except Exception as e:
    print(f"Ошибка инициализации InferenceClient: {e}. Проверьте токен и доступность модели.")

# --- Определение инструментов --- 

def calculator(expression: str) -> str:
    """Выполняет математические вычисления."""
    try:
        # Ограничение на доступные функции для безопасности
        allowed_names = {k: v for k, v in math.__dict__.items() if not k.startswith("__")}
        allowed_names["abs"] = abs
        allowed_names["round"] = round
        allowed_names["max"] = max
        allowed_names["min"] = min
        
        # Удаление потенциально опасных символов
        safe_expression = re.sub(r"[^0-9\.\+\-\*\/\(\)\s]|\b(import|exec|eval|open|lambda|\_\_)\b", "", expression)
        
        if safe_expression != expression:
            return "Ошибка: Обнаружены недопустимые символы в выражении."
            
        result = eval(safe_expression, {"__builtins__": {}}, allowed_names)
        return f"Результат: {result}"
    except Exception as e:
        return f"Ошибка в вычислении: {str(e)}"

def web_search(query: str) -> str:
    """Выполняет поиск в интернете по заданному запросу (симуляция)."""
    # Простая симуляция для теста
    if "погода" in query.lower():
        return "В городе, который вы ищете, сегодня солнечно, +25C."
    elif "hugging face" in query.lower():
        return "Hugging Face - это платформа и сообщество для работы с моделями машинного обучения."
    elif "python" in query.lower():
        return "Python - высокоуровневый язык программирования общего назначения, созданный Гвидо ван Россумом."
    else:
        return f"По вашему запросу '{query}' найдена общая информация."

# --- Функция для запуска агента --- 
def run_agent(query: str) -> str:
    """Запускает агента для ответа на вопрос."""
    if client is None:
        return "Ошибка: Клиент Hugging Face не инициализирован. Проверьте токен и доступность модели."
    
    # Определение инструментов для модели
    tools = [
        {
            "type": "function",
            "function": {
                "name": "calculator",
                "description": "Выполняет математические вычисления",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "expression": {
                            "type": "string",
                            "description": "Математическое выражение для вычисления"
                        }
                    },
                    "required": ["expression"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "web_search",
                "description": "Ищет информацию в интернете",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Поисковый запрос"
                        }
                    },
                    "required": ["query"]
                }
            }
        }
    ]
    
    # Начальное сообщение от пользователя
    messages = [{"role": "user", "content": query}]
    
    # Максимальное количество итераций
    max_iterations = 5
    
    for _ in range(max_iterations):
        # Вызов модели
        response = client.chat_completion(
            messages=messages,
            tools=tools,
            tool_choice="auto",
            temperature=0.1,
            max_tokens=1024
        )
        
        # Получение ответа модели
        assistant_message = response["choices"][0]["message"]
        messages.append(assistant_message)
        
        # Проверка на наличие вызовов инструментов
        if "tool_calls" in assistant_message and assistant_message["tool_calls"]:
            for tool_call in assistant_message["tool_calls"]:
                # Получение имени и аргументов инструмента
                function_name = tool_call["function"]["name"]
                function_args = json.loads(tool_call["function"]["arguments"])
                
                # Вызов соответствующего инструмента
                if function_name == "calculator":
                    result = calculator(function_args["expression"])
                elif function_name == "web_search":
                    result = web_search(function_args["query"])
                else:
                    result = f"Инструмент {function_name} не найден."
                
                # Добавление результата в сообщения
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call["id"],
                    "name": function_name,
                    "content": result
                })
        else:
            # Если нет вызовов инструментов, возвращаем ответ
            return assistant_message["content"]
    
    # Если достигнуто максимальное количество итераций, возвращаем последний ответ
    return messages[-1]["content"]

# --- Создание Gradio интерфейса ---
def gradio_interface(query):
    """Обработчик для Gradio интерфейса."""
    if not query.strip():
        return "Пожалуйста, введите вопрос."
    
    try:
        response = run_agent(query)
        return response
    except Exception as e:
        return f"Произошла ошибка: {str(e)}"

# Создание интерфейса
demo = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=2, placeholder="Введите ваш вопрос здесь..."),
    outputs="text",
    title="Agent Challenge - Финальный агент",
    description="Этот агент может отвечать на вопросы, выполнять математические вычисления и искать информацию."
)

# Запуск интерфейса
if __name__ == "__main__":
    demo.launch()