from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

@tool
def keyword_research_tool(topic: str, max_results: int) -> str:
    """Генерирует релевантные ключевые слова для заданной темы
    Args:
        topic: Основная тема для генерации ключевых слов
        max_results: Максимальное количество возвращаемых ключевых слов
    """
    import random
    from typing import List
    
    # Базовые модификаторы для расширения темы
    modifiers = [
        "лучшие", "топ", "как", "советы", "рецепты",
        "польза", "вред", "отзывы", "купить", "сделать",
        "обзор", "курс", "уроки", "для начинающих", 
        "2023", "идеи", "примеры", "методы", "технологии",
        "тренды", "планы", "стоимость", "бесплатно", "преимущества"
    ]
    
    # Разбиваем тему на основные слова
    seed_words = [word.strip() for word in topic.split() if len(word) > 3]
    
    # Генерируем комбинации
    candidates = set()
    
    # Добавляем базовые варианты
    candidates.add(topic)
    candidates.add(f"лучшие {topic}")
    candidates.add(f"{topic} 2023")
    
    # Генерируем комбинации с модификаторами
    for modifier in modifiers:
        candidates.add(f"{modifier} {topic}")
        candidates.add(f"{topic} {modifier}")
        for word in seed_words:
            candidates.add(f"{modifier} {word}")
            candidates.add(f"{word} {modifier}")
    
    # Добавляем вопросы
    question_words = ["как", "что", "где", "когда", "почему"]
    for qw in question_words:
        candidates.add(f"{qw} {topic}")
        for word in seed_words:
            candidates.add(f"{qw} {word}")
    
    # Преобразуем в список и перемешиваем
    results = list(candidates)
    random.shuffle(results)
    
    # Выбираем запрошенное количество
    return ", ".join(results[:max(max_results, 1)])

# Initialize components
final_answer = FinalAnswerTool()

# Model initialization
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

# Load tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Load prompt templates
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Create agent
agent = CodeAgent(
    tools=[
        image_generation_tool,
        keyword_research_tool,
        final_answer,
        DuckDuckGoSearchTool()
    ],
    model=model,
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Launch UI
GradioUI(agent).launch()