from smolagents import CodeAgent, HfApiModel, load_tool, tool, DuckDuckGoSearchTool
import datetime
import requests
import pytz
import yaml
import base64
import os
import io
from PIL import Image
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# 1. Сначала определяем ВСЕ функции-инструменты

@tool
def web_search(query: str) -> str:
    """Search the web for information using DuckDuckGo.
    Args:
        query: The search query to look up
    """
    try:
        search_tool = DuckDuckGoSearchTool()
        results = search_tool(query)
        return f"🔍 **Результаты поиска по '{query}':**\n\n{results}"
    except Exception as e:
        return f"❌ Ошибка поиска: {str(e)}"

@tool
def get_weather(city: str) -> str:
    """Get current weather for a city using wttr.in service.
    Args:
        city: The name of the city to get weather for (e.g., 'Moscow', 'London')
    """
    try:
        # Используем wttr.in с метрической системой (Цельсий)
        url = f"http://wttr.in/{city}?format=%C+%t+%h+%w+%P&lang=ru&m"
        response = requests.get(url, timeout=10)
        
        if response.status_code == 200:
            data = response.text.strip()
            
            # Парсим ответ wttr.in
            parts = data.split(' ')
            if len(parts) >= 4:
                condition = parts[0]  # Погодные условия
                temperature = parts[1]  # Температура
                humidity = parts[2]    # Влажность
                wind = parts[3]        # Ветер
                
                return (f"🌤 Погода в {city}:\n"
                       f"🌡 {temperature}\n"
                       f"☁️ {condition}\n"ы
                       f"💧 Влажность: {humidity}\n"
                       f"💨 Ветер: {wind}")
            else:
                return f"Погода в {city}: {data}"
        else:
            return f"❌ Не удалось получить погоду для {city}"
            
    except Exception as e:
        return f"❌ Ошибка: {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York')
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

@tool
def generate_image(prompt: str) -> str:
    """Generate an image from text prompt and save it.
    Args:
        prompt: Text description of the image to generate
    """
    try:
        image = image_generation_tool(prompt=prompt)
        
        # Создаем папку для изображений если её нет
        os.makedirs("images", exist_ok=True)
        
        # Сохраняем в папку images
        import time
        filename = f"image_{int(time.time())}.png"
        filepath = f"images/{filename}"
        image.save(filepath)
        
        return f"🎨 **Изображение сгенерировано!**\n\n" \
               f"**Запрос:** {prompt}\n" \
               f"**Файл:** `{filepath}`\n\n" \
               f"📁 Изображение сохранено в папке 'images/'"
               
    except Exception as e:
        return f"❌ Ошибка генерации: {str(e)}"

@tool
def calculate_math(expression: str) -> str:
    """Calculate mathematical expressions safely.
    Args:
        expression: A mathematical expression to calculate (e.g., '2+2', '5*3/2')
    """
    try:
        # Безопасные математические операции
        allowed_chars = set('0123456789+-*/.() ')
        if all(c in allowed_chars for c in expression):
            result = eval(expression)
            return f"🧮 Результат: {result}"
        else:
            return "❌ Используйте только цифры и +-*/.()"
    except:
        return "❌ Ошибка в выражении"


@tool
def coin_flip() -> str:
    """Flip a coin and get heads or tails."""
    import random
    result = random.choice(["орел", "решка"])
    return f"🪙"решка"])
    return f"🪙 Результат броска: **{result}**"

@tool
def generate_password(length: int = 12) -> str:
    """Generate a secure random password.
    Args:
        length: The length of the password (default: 12)
    """
    import string
    import random
    characters = string.ascii_letters + string.digits + "!@#$%"
    password = ''.join(random.choice(characters) for _ in range(length))
    return f"🔐 Сгенерированный пароль: `{password}`"

@tool
def advice_by_category(category: str = "random") -> str:
    """Get advice by specific category.
    Args:
        category: Advice category ('motivation', 'health', 'work', 'relationships', 'random', 'personal development')
    """
    advice_categories = {
        'motivation': [
            "Ты сильнее, чем думаешь 💪",
            "Каждый эксперт когда-то был новичком 🌱",
            "Не смотри на других - соревнуйся с собой вчерашним 🏆",
            "Дыши глубоко в стрессовых ситуациях 🧘",
            "Помни: все временно, и это тоже пройдет 🌈",
            "Прими то, что не можешь изменить 🌊",
            "Найди время для тишины каждый день 🤫",
        ],
        'health': [
            "Прогулка на свежем воздухе творит чудеса 🌳",
            "Слушай свое тело - оно мудрее любого доктора 👂",
            "Здоровье - главное богатство 💎",
            "Пей больше воды 💧",
            "Двигайся каждый день 🏃‍♂️",
            "Высыпайся - сон это суперсила 😴",
            "Ешь больше овощей и фруктов 🍎",
        ],
        'work': [
            "Делай зарядку с утра 🐸",
            "Учись делегировать 🤲",
            "Баланс работы и отдыха - ключ к успеху ⚖️",
            "Инвестируй в свои знания 💼",
            "Откладывай хотя бы 10% от дохода 💰",
            "Учись говорить 'нет' когда нужно 🚫",
            "Цени свой труд и время ⭐",
        ],
        'relationships': [
            "Искренний комплимент может изменить чей-то день 🌈",
            "Слушай чтобы понять, а не чтобы ответить 💭",
            "Прощение освобождает в первую очередь тебя 🕊️",
            "Будь добрым к себе и другим 💖",
            "Цени моменты - они уникальны 🌟",
            "Говори 'спасибо' чаще 🙏",
            "Слушай больше, чем говоришь 👂"
        ],
        'personal development': [
            "Выходи из зоны комфорта каждый день 🚪",
            "Учись новому навыку каждый месяц 🎓",
            "Рефлексируй над своим днем 📝",
            "Ставь реалистичные цели 🎯",
        ],
    }
    
    import random
    
    if category == 'random' or category not in advice_categories:
        all_advice = [advice for category_list in advice_categories.values() for advice in category_list]
        return f"💡 Совет: {random.choice(all_advice)}"
    else:
        return f"💡 {category.title()} совет: {random.choice(advice_categories[category])}"

@tool
def calculate_age(birth_year: int) -> str:
    """Calculate current age based on birth year.
    Args:
        birth_year: The year of birth
    """
    from datetime import datetime
    current_year = datetime.now().year
    age = current_year - birth_year
    
    if age < 0:
        return "❌ Год рождения не может быть в будущем!"
    elif age == 0:
        return "👶 0 лет! Новое начало! 🎉"
    elif age < 13:
        return f"🎂 {age} лет! Детство - прекрасная пора! 🌈"
    elif age < 20:
        return f"🎂 {age} лет! Юность полна возможностей! 🚀"
    elif age < 30:
        return f"🎂 {age} лет! Молодость и энергия! ⚡"
    elif age < 40:
        return f"🎂 {age} лет! Опыт и уверенность! 💼"
    elif age < 50:
        return f"🎂 {age} лет! Расцвет сил и мудрости! 🌟"
    elif age < 60:
        return f"🎂 {age} лет! Золотой возраст! 🏆"
    elif age < 120:
        return f"🎂 {age} лет! Богатый жизненный опыт! 📚"
    else:
        return f"🎂 {age} лет! Легендарное долголетие! 🎊"




# 2. ТОЛЬКО ПОСЛЕ ВСЕХ ФУНКЦИЙ создаем остальные объекты

final_answer = FinalAnswerTool()

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

# Загрузка инструмента генерации изображений
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

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

# Добавляем проверку для final_answer
if "final_answer" not in prompt_templates:
    prompt_templates["final_answer"] = {
        "pre_messages": "Based on my research: ",
        "post_messages": ""
    }
    
# 3. Теперь используем функции в списке инструментов
agent = CodeAgent(
    model=model,
    tools=[final_answer, web_search, generate_image, get_current_time_in_timezone, get_weather, calculate_math, coin_flip, generate_password, advice_by_category, calculate_age],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()