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

# -------------------------
# 1) Real custom tool: translate text via LibreTranslate API
# -------------------------
@tool
def translate_text(text: str, target_lang: str) -> str:
    """Translate a given text into the target language using LibreTranslate.
    Args:
        text: the source text to translate
        target_lang: ISO code of the target language (e.g. 'it', 'en', 'es')
    """
    try:
        response = requests.post(
            'https://libretranslate.com/translate',
            json={
                'q': text,
                'source': 'auto',
                'target': target_lang,
                'format': 'text'
            },
            timeout=5
        )
        data = response.json()
        return data.get('translatedText', 'Errore nella traduzione')
    except Exception as e:
        return f"Translation error: {e}"

# -------------------------
# 2) Prebuilt tool: current time in timezone
# -------------------------
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Fetch the current local time in a specified timezone.
    Args:
        timezone: e.g., 'Europe/Rome', 'America/New_York'
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Current local time in {timezone}: {local_time}"
    except Exception as e:
        return f"Error fetching time: {e}"

# -------------------------
# 3) Load additional tools
# -------------------------
# DuckDuckGo web search tool
duck_search = DuckDuckGoSearchTool()

# Image generation tool from Hub
image_gen_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Final answer formatter
tool_final = FinalAnswerTool()

# -------------------------
# 4) Model setup
# -------------------------
model = HfApiModel(
    max_tokens=1500,
    temperature=0.7,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct'
)

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

# -------------------------
# 6) Agent construction
# -------------------------
agent = CodeAgent(
    model=model,
    tools=[
        tool_final,
        translate_text,
        get_current_time_in_timezone,
        duck_search,
        image_gen_tool,
    ],
    max_steps=8,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

# -------------------------
# 7) Launch Gradio UI
# -------------------------
if __name__ == "__main__":
    GradioUI(agent).launch()