from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import pytz
import yaml

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


@tool
def smart_utility_tool(operation: str, value1: float, value2: float = 0) -> str:
    """
    A smart utility tool that performs basic mathematical operations.

    Args:
        operation: The type of operation to perform.
                   Options: 'add', 'subtract', 'multiply', 'divide', 'power'
        value1: The first numeric value
        value2: The second numeric value (optional for some operations)

    Returns:
        A string containing the result of the operation.
    """
    try:
        if operation == "add":
            result = value1 + value2
        elif operation == "subtract":
            result = value1 - value2
        elif operation == "multiply":
            result = value1 * value2
        elif operation == "divide":
            if value2 == 0:
                return "Error: Cannot divide by zero."
            result = value1 / value2
        elif operation == "power":
            result = value1 ** value2
        else:
            return (
                f"Unsupported operation '{operation}'. "
                "Choose from: add, subtract, multiply, divide, power."
            )

        return f"Result: {operation}({value1}, {value2}) = {result}"
    except Exception as e:
        return f"Error performing operation: {str(e)}"


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Fetch current local time in a specified timezone.

    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').

    Returns:
        A human-readable string with the current time in that 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)}"


# --- Final answer tool (keep this) ---
final_answer = FinalAnswerTool()

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

# --- Pre-made tools ---
web_search = DuckDuckGoSearchTool()

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

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

# --- Agent ---
agent = CodeAgent(
    model=model,
    tools=[
        final_answer,              # don't remove
        web_search,                # premade tool
        image_generation_tool,     # hub tool
        get_current_time_in_timezone,
        smart_utility_tool,
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
)

# --- UI ---
GradioUI(agent).launch()