from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# --- 1. Islamic Info Tool ---
@tool
def get_islamic_prayer_times(city: str, country: str) -> str:
    """Fetches the Islamic prayer times for a specific city and country.
    Args:
        city: The name of the city (e.g., 'London').
        country: The name of the country (e.g., 'UK').
    """
    try:
        # Using the Aladhan API to get timings
        url = f"http://api.aladhan.com/v1/timingsByCity?city={city}&country={country}&method=2"
        response = requests.get(url)
        data = response.json()
        if data['code'] == 200:
            timings = data['data']['timings']
            date = data['data']['date']['readable']
            return (f"Prayer times for {city}, {country} on {date}:\n"
                    f"Fajr: {timings['Fajr']}, Dhuhr: {timings['Dhuhr']}, "
                    f"Asr: {timings['Asr']}, Maghrib: {timings['Maghrib']}, Isha: {timings['Isha']}")
        else:
            return f"Could not find prayer times for {city}. Please check the spelling."
    except Exception as e:
        return f"Error fetching prayer times: {str(e)}"

# --- 2. Timezone Tool ---
@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 Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

# --- 3. Placeholder Tool (Customizable) ---
@tool
def my_custom_tool(arg1:str, arg2:int)-> str:
    """A tool that currently returns a greeting.
    Args:
        arg1: A name or string.
        arg2: An integer.
    """
    return f"Hello {arg1}, your magic number is {arg2}!"

# --- 4. Setup Model & Tools ---
final_answer = FinalAnswerTool()

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

model = InferenceClientModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

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

# --- 5. Initialize Agent ---
agent = CodeAgent(
    model=model,
    tools=[
        final_answer, 
        get_current_time_in_timezone, 
        get_islamic_prayer_times, 
        image_generation_tool, 
        my_custom_tool
    ], 
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="SmartAssistant",
    description="An agent that can tell time, provide prayer schedules, and generate art.",
    prompt_templates=prompt_templates
)

# --- 6. Launch UI ---
if __name__ == "__main__":
    GradioUI(agent).launch()