from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
# from tools.final_answer import FinalAnswerTool  # Assuming this is in a 'tools' directory
from tools.final_answer import FinalAnswerTool  # <--- Correct relative import

# For Gradio (assuming you have a GradioUI class defined)
from Gradio_UI import GradioUI  # Or wherever you have GradioUI defined
import os

# Set your Hugging Face API token (replace with your actual token)
# It's BEST to set this as an environment variable:  # Replace!


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    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 pytz.exceptions.UnknownTimeZoneError:
        return f"Error: Unknown timezone '{timezone}'. Please provide a valid timezone."
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

# Use a valid and accessible Llama 3 model ID
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='meta-llama/Llama-3.2-1B-Instruct',  # Corrected Model ID
    # api_key="hf_YOUR_API_KEY",  # Use environment variable instead
    custom_role_conversions={  # Correct role conversions for Llama 3
        "system": "<|start_header_id|>system<|end_header_id|>\n\n{content}<|eot_id|>",
        "user": "<|start_header_id|>user<|end_header_id|>\n\n{content}<|eot_id|>",
        "assistant": "<|start_header_id|>assistant<|end_header_id|>\n\n{content}<|eot_id|>",
    },
)


# --- Safer Image Handling (DuckDuckGo Search) ---
ddg_search = DuckDuckGoSearchTool()



with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone, ddg_search],  # Use DDG instead of untrusted remote code
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="Time and Image Agent",  # Give your agent a name
    description="An agent that can tell you the time and search for images.",  # And a description
    prompt_templates=prompt_templates,
)


# --- Launch Gradio UI ---
GradioUI(agent).launch(debug=True)  # Add debug=True for easier debugging