from smolagents import CodeAgent, DuckDuckGoSearchTool, tool
import datetime
import requests
import pytz
import yaml

# Создаем свой собственный инструмент для финального ответа
@tool
def final_answer(answer: str) -> str:
    """Provide the final answer to the user. Use this when you have the answer to the user's question.
    
    Args:
        answer: The final answer to present to the user
    """
    return answer

# Инструменты (tools), используемые агентом:
# - web_search: поиск в интернете через DuckDuckGo
# - get_weather: получение текущей погоды для города
# - get_weather_detailed: подробный прогноз погоды
# - convert_fahrenheit_to_celsius: конвертация температуры
# - get_current_time_in_timezone: время в указанном часовом поясе
# - calculate_math: вычисление математических выражений
# - suggest_weather_sources: рекомендации источников погоды
# - final_answer: инструмент для финального ответа

@tool
def web_search(query: str) -> str:
    """Search the web for information using DuckDuckGo."""
    try:
        search_tool = DuckDuckGoSearchTool()
        results = search_tool(query)
        if not results or "No results found" in results:
            return f"Поиск не дал результатов для запроса: '{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 reliable weather services."""
    try:
        url = f"https://wttr.in/{city}?format=%C+%t+%h+%w&m&lang=ru"
        response = requests.get(url, timeout=10)
        
        if response.status_code == 200 and response.text.strip():
            data = response.text.strip()
            if '°F' in data:
                data = data.replace('°F', '°C')
            return f"Погода в {city}:\n{data}"
        
        url2 = f"https://wttr.in/{city}?format=%C+%t+%h+%w+%P&m&lang=ru"
        response2 = requests.get(url2, timeout=10)
        
        if response2.status_code == 200 and response2.text.strip():
            data = response2.text.strip()
            if '°F' in data:
                data = data.replace('°F', '°C')
            return f"Погода в {city}:\n{data}"
            
        return f"Данные о погоде в {city} временно недоступны. Попробуйте позже или проверьте на сайте gismeteo.ru"
            
    except Exception as e:
        return f"Ошибка при получении погоды: {str(e)}"

@tool
def get_weather_detailed(city: str) -> str:
    """Get detailed weather forecast for a city."""
    try:
        url = f"https://wttr.in/{city}?m&lang=ru"
        response = requests.get(url, timeout=10)
        
        if response.status_code == 200:
            lines = response.text.split('\n')
            short_forecast = '\n'.join(lines[:8])
            short_forecast = short_forecast.replace('°F', '°C')
            return f"Подробный прогноз для {city}:\n{short_forecast}"
        else:
            return f"Не удалось получить подробный прогноз для {city}"
            
    except Exception as e:
        return f"Ошибка: {str(e)}"

@tool
def convert_fahrenheit_to_celsius(f_temp: float) -> str:
    """Convert Fahrenheit temperature to Celsius."""
    try:
        c_temp = (f_temp - 32) * 5/9
        return f"{f_temp}°F = {c_temp:.1f}°C"
    except Exception as e:
        return f"Ошибка конвертации: {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Get current local time in specified timezone."""
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Текущее время в {timezone}: {local_time}"
    except Exception as e:
        return f"Ошибка получения времени для часового пояса '{timezone}': {str(e)}"

@tool
def calculate_math(expression: str) -> str:
    """Calculate mathematical expressions safely."""
    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 suggest_weather_sources(city: str) -> str:
    """Suggest reliable weather sources for a city."""
    sources = [
        f"Gismeteo: https://www.gismeteo.ru/weather-{city.lower()}-4368/",
        f"Yandex.Погода: https://yandex.ru/pogoda/{city.lower()}",
        f"Weather.com: https://weather.com/weather/today/l/{city}",
        f"AccuWeather: https://www.accuweather.com/ru/ru/{city.lower()}/weather-forecast"
    ]
    
    return f"Надежные источники погоды для {city}:\n" + "\n".join(sources)

# Используем локальную модель вместо HfApiModel
try:
    from smolagents import TransformersModel
    
    model = TransformersModel(
        model_id='Qwen/Qwen2.5-Coder-1.5B',
        max_tokens=2096,
        temperature=0.5
    )
except ImportError:
    # Если TransformersModel недоступен, используем LiteLLM
    from smolagents import LiteLLMModel
    model = LiteLLMModel(
        model_id='gpt-3.5-turbo',
        max_tokens=2096,
        temperature=0.5
    )

# Загружаем промпты, если файл существует, иначе создаем базовый шаблон
try:
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
except FileNotFoundError:
    prompt_templates = {
        "final_answer": {
            "pre_messages": "Based on my research: ",
            "post_messages": ""
        }
    }

if "final_answer" not in prompt_templates:
    prompt_templates["final_answer"] = {
        "pre_messages": "Based on my research: ",
        "post_messages": ""
    }
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, web_search, get_weather, get_weather_detailed, convert_fahrenheit_to_celsius, get_current_time_in_timezone, calculate_math, suggest_weather_sources],
    max_steps=8,
    verbosity_level=0,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Простой консольный интерфейс вместо GradioUI
def main():
    print("🤖 ИИ-агент готов к работе!")
    print("Доступные функции: поиск в интернете, погода, время, математические вычисления")
    print("Введите 'выход' или 'exit' для завершения работы\n")
    
    while True:
        try:
            user_input = input("Ваш вопрос: ").strip()
            
            if user_input.lower() in ['выход', 'exit', 'quit']:
                print("До свидания!")
                break
                
            if not user_input:
                continue
                
            print("🤔 Думаю...")
            response = agent.run(user_input)
            print(f"🤖 Ответ: {response}\n")
            
        except KeyboardInterrupt:
            print("\nДо свидания!")
            break
        except Exception as e:
            print(f"❌ Произошла ошибка: {e}\n")

if __name__ == "__main__":
    main()