from smolagents import CodeAgent, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import os
from dotenv import load_dotenv
from tools.web_search import DuckDuckGoSearchTool
from tools.visit_webpage import VisitWebpageTool
from tools.final_answer import final_answer
import gradio as gr

load_dotenv()

@tool
def get_weather(city: str) -> str:
    """Get current weather conditions for a specified city
    
    Args:
        city (str): The name of the city to check weather for
    """
    try:
        response = requests.get(f"https://wttr.in/{city}?format=%C+%t")
        response.raise_for_status()
        return f"Weather in {city}: {response.text}"
    except Exception as e:
        return f"Weather check failed: {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Get current local time in specified timezone
    
    Args:
        timezone (str): A valid timezone identifier
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Current time in {timezone}: {local_time}"
    except Exception as e:
        return f"Error: {str(e)}"

# Model configuration - minimal and working
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") as f:
    prompt_templates = yaml.safe_load(f)
    
agent = CodeAgent(
    model=model,
    tools=[
        get_weather,
        get_current_time_in_timezone,
        image_generation_tool,
        DuckDuckGoSearchTool(),
        VisitWebpageTool(),
        final_answer
    ],
    max_steps=10,
    verbosity_level=2,
    prompt_templates=prompt_templates
)

def run_agent(query: str) -> str:
    """Wrapper function for Gradio interface"""
    try:
        return str(agent.run(query))
    except Exception as e:
        return f"Error: {str(e)}"

if __name__ == "__main__":
    if not os.getenv("HF_TOKEN"):
        raise ValueError("HF_TOKEN environment variable not set")
    
    gr.Interface(
        fn=run_agent,
        inputs=gr.Textbox(label="Input"),
        outputs=gr.Textbox(label="Output"),
        title="AI Agent"
    ).launch(server_name="0.0.0.0", share=True)