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 generate_keywords(topic: str, max_keywords: int = 30) -> str:
    """Генерирует релевантные ключевые слова для заданной темы
    Args:
        topic: Тема для генерации (на русском языке)
        max_keywords: Максимальное количество возвращаемых ключевых слов
    """
    import random
    from collections import deque
    
    # Инициализация компонентов генерации
    base_words = [w.strip().lower() for w in topic.split() if len(w) > 2]
    if not base_words:
        return "Необходимо указать более конкретную тему"
    
    # Паттерны для генерации
    patterns = [
        ["вопросы", "что такое", "как сделать", "почему", "когда", "где найти"],
        ["топ", "лучшие", "рейтинг", "самые популярные"],
        ["методы", "способы", "технологии", "решения"],
        ["обзор", "отзывы", "сравнение", "виды"],
        ["курс", "урок", "обучение", "самоучитель"],
        ["2025", "новинки", "тренды", "будущее"],
        ["для начинающих", "профессиональные", "премиум"],
        ["бесплатно", "купить", "заказать", "стоимость"],
        ["преимущества", "недостатки", "польза", "вред"],
        ["идеи", "примеры", "кейсы", "истории успеха"]
    ]
    
    # Генерация вариантов
    keywords = set()
    queue = deque([(0, [], base_words)])
    
    while queue:
        level, current, remaining = queue.popleft()
        
        if level >= 3 or not remaining:
            if current:
                keywords.add(" ".join(current))
            continue
            
        # Добавляем паттерны
        if level == 0:
            for pattern_group in patterns:
                for p in pattern_group:
                    new_phrase = [p] + current
                    queue.append((level+1, new_phrase, remaining))
        
        # Комбинируем слова
        for i in range(len(remaining)):
            next_word = remaining[i]
            new_phrase = current + [next_word]
            new_remaining = remaining[:i] + remaining[i+1:]
            queue.append((level+1, new_phrase, new_remaining))
            
            # Добавляем модификаторы
            modifiers = ["", "2025", "курс", "обзор", "топ", "способы"]
            for mod in modifiers:
                if mod:
                    modified = current + [f"{next_word}-{mod}"]
                    queue.append((level+2, modified, new_remaining))
    
    # Ранжирование и выборка
    result = sorted(keywords, 
                   key=lambda x: (-len(x.split()), x))
    random.shuffle(result[:len(result)//2])
    
    return ", ".join(result[:max(max_keywords, 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,
        generate_keywords,
        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()