import datetime
from pathlib import Path

import pytz
import yaml

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

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


# 1. Custom tool
@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """
    Demonstration tool that combines two inputs.

    Args:
        arg1: A text message.
        arg2: An integer value.
    """
    return f"Message: {arg1}. Number received: {arg2}."


# 2. Time-zone tool
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Get current local time in a specified timezone.

    Args:
        timezone: Timezone such as Asia/Riyadh.
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz)

        return (
            f"Current time in {timezone}: "
            f"{local_time.strftime('%Y-%m-%d %H:%M:%S')}"
        )

    except pytz.exceptions.UnknownTimeZoneError:
        return f"Unknown timezone: {timezone}"


# 3. Final answer tool
final_answer = FinalAnswerTool()


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


# 5. Load image-generation tool
image_generation_tool = load_tool(
    "agents-course/text-to-image",
    trust_remote_code=True,
)


# 6. Load prompts
BASE_DIR = Path(__file__).resolve().parent

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


# 7. Create agent
agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        my_custom_tool,
        get_current_time_in_timezone,
        DuckDuckGoSearchTool(),
        image_generation_tool,
    ],
    max_steps=8,
    verbosity_level=1,
    planning_interval=None,
    name="First Agent",
    description="A general-purpose AI agent with search, time and image tools.",
    prompt_templates=prompt_templates,
)


# 8. Launch application
if __name__ == "__main__":
    GradioUI(agent).launch()