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

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@tool
def my_custom_tool(arg1: int, arg2: int) -> int:
    """Add two numbers.

    Args:
        arg1: The first number.
        arg2: The second number.
    """
    return arg1 + arg2


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

    Args:
        timezone: A valid timezone, such as Asia/Dubai or America/New_York.
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime(
            "%Y-%m-%d %H:%M:%S %Z"
        )
        return f"The current local time in {timezone} is: {local_time}"

    except pytz.exceptions.UnknownTimeZoneError:
        return (
            f"Invalid timezone: {timezone}. "
            "Try a valid timezone such as Asia/Dubai."
        )
    except Exception as exc:
        logger.exception("Failed to fetch timezone")
        return f"Error fetching time: {exc}"


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

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

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

    if not isinstance(prompt_templates, dict):
        raise ValueError("prompts.yaml must contain a YAML mapping.")

    final_answer = FinalAnswerTool()

    # Initialize the agent with all tools
    agent = CodeAgent(
        model=model,
        tools=[
            get_current_time_in_timezone,
            my_custom_tool,
            image_generation_tool,
            final_answer,
        ],
        max_steps=6,
        verbosity_level=1,
        grammar=None,
        planning_interval=None,
        name="AI Assistant",
        description=(
            "An AI assistant that can calculate, retrieve local time, "
            "and generate images."
        ),
        prompt_templates=prompt_templates,
    )

    # Do not pass share=True here; your Gradio_UI.py
    # already sets share=True internally.
    GradioUI(agent).launch()


if __name__ == "__main__":
    main()