from smolagents import (
    CodeAgent,
    HfApiModel,
    load_tool,
    tool,
)

import datetime
import pytz
import re

from tools.final_answer import FinalAnswerTool
from tools.web_search import DuckDuckGoSearchTool
from tools.visit_webpage import VisitWebpageTool

from Gradio_UI import GradioUI


# ==========================================================
# CUSTOM TOOLS
# ==========================================================

@tool
def calculate_min_price(prices: list[float]) -> str:
    """
    Calculates the minimum value from a list of prices.

    Args:
        prices: List of prices.
    """
    return f"The minimum price is {min(prices)}"


@tool
def extract_price_from_snippet(snippet: str) -> list[str]:
    """
    Extracts prices from a block of text.

    Args:
        snippet: Text containing prices.
    """

    pattern = (
        r"\$\d+(?:,\d{3})*(?:\.\d{2})?"
        r"|\d+(?:,\d{3})*(?:\.\d{2})?\s*"
        r"(?:USD|EUR|GBP|INR|AUD|CAD)?"
    )

    return re.findall(pattern, snippet)


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Returns current local time.

    Args:
        timezone: Valid pytz timezone.
    """

    try:
        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz)
        return now.strftime("%Y-%m-%d %H:%M:%S")

    except Exception as e:
        return str(e)


# ==========================================================
# REQUIRED FINAL ANSWER TOOL
# ==========================================================

final_answer = FinalAnswerTool()


# ==========================================================
# WEB TOOLS
# ==========================================================

web_search = DuckDuckGoSearchTool()
visit_webpage = VisitWebpageTool()


# ==========================================================
# OPTIONAL IMAGE TOOL
# ==========================================================

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


# ==========================================================
# MODEL
# ==========================================================

import os

model = HfApiModel(
    model_id="Qwen/Qwen2.5-72B-Instruct",
    token=os.environ["HF_TOKEN"],
    max_tokens=2048,
    temperature=0.1,
)


# ==========================================================
# AGENT
# ==========================================================

agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        web_search,
        visit_webpage,
        calculate_min_price,
        extract_price_from_snippet,
        get_current_time_in_timezone,
    ],
    max_steps=12,
    verbosity_level=2,
    planning_interval=2,
)


# ==========================================================
# UI
# ==========================================================

GradioUI(agent).launch()