# app.py

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    load_tool,
    tool,
    HfApiModel,   # replacement for InferenceClientModel
)
import datetime
import re
import requests
import pytz
import yaml
from typing import Annotated

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


# ---- Custom Tools ----

@tool
def get_current_time_in_timezone(
    timezone: Annotated[str, "IANA timezone like 'America/Los_Angeles'"]
) -> str:
    """Return the current local time for the given timezone.

    Args:
        timezone (str): IANA timezone name, e.g., 'America/Los_Angeles' or 'Asia/Tokyo'.

    Returns:
        str: A human-readable string with the local time in the requested timezone.
    """
    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 fetch_page_title(
    url: Annotated[str, "An http/https URL to fetch"]
) -> str:
    """Fetch the page and return its <title> text.

    Args:
        url (str): A valid HTTP/HTTPS URL to retrieve.

    Returns:
        str: The page title, or an error message if fetching/parsing fails.
    """
    try:
        resp = requests.get(url, timeout=10)
        resp.raise_for_status()
        m = re.search(r"<title[^>]*>(.*?)</title>", resp.text, flags=re.IGNORECASE | re.DOTALL)
        title = re.sub(r"\s+", " ", m.group(1)).strip() if m else "(no <title> found)"
        return f"Title: {title}"
    except Exception as e:
        return f"Error fetching title from '{url}': {e}"


@tool
def echo_and_length(
    text: Annotated[str, "Any text to echo back"]
) -> str:
    """Echo the text and report its character length.

    Args:
        text (str): The text to echo back.

    Returns:
        str: The echoed text and its character count.
    """
    t = text.strip()
    return f"Echo: {t}\nLength: {len(t)}"


# ---- Model & Tools wiring ----

final_answer = FinalAnswerTool()

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

image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
web_search = DuckDuckGoSearchTool()

with open("prompts.yaml", "r", encoding="utf-8") as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        web_search,
        image_generation_tool,
        get_current_time_in_timezone,
        fetch_page_title,
        echo_and_length,
    ],
    max_steps=8,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="my_first_smolagents_agent",
    description="A tiny agent that can search, fetch titles, tell time, and generate images.",
    prompt_templates=prompt_templates,
)

GradioUI(agent).launch()