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

# --- Define Tools ---

@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """
    A placeholder tool.
    
    Args:
        arg1: The first argument.
        arg2: The second argument.
    
    Returns:
        A placeholder string.
    """
    return "What magic will you build?"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Fetches the current local time in a specified timezone and adds a friendly, reflective comment.
    
    Args:
        timezone: A valid timezone string (e.g., 'America/New_York').
    
    Returns:
        A string with the current time and a comment.
    """
    try:
        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz)
        local_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
        hour = now.hour
        if hour < 6 or hour >= 22:
            comment = "🌙 It's late, and the night whispers secrets. What dreams keep you awake?"
        else:
            comment = "☀️ The day unfolds with gentle light. Where are you enjoying this beautiful moment?"
        return f"The current local time in {timezone} is: {local_time_str}. {comment}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

@tool
def get_weather(location: str) -> str:
    """
    Simulates fetching poetic weather information for a given location.
    
    Args:
        location: The name of your location.
    
    Returns:
        A poetic description of the current weather.
    """
    # In a real implementation, this tool would call a weather API.
    # For now, we'll return a simulated weather description.
    return f"In {location}, the weather is radiant and gentle, with soft breezes and a sky that whispers tales of wonder."

@tool
def probe_story_preferences(likes: str, dislikes: str, favorite_theme: str) -> str:
    """
    Captures your story preferences and returns a poetic summary.
    
    Args:
        likes: What you cherish in stories.
        dislikes: What you prefer to leave behind.
        favorite_theme: Your favored narrative theme (e.g., adventure, mystery).
    
    Returns:
        A summary that helps shape your personalized narrative.
    """
    return f"Splendid! You delight in {likes} 🌸, wish to steer clear of {dislikes} 🚫, and are enchanted by {favorite_theme} themes. Let us craft your story accordingly."

@tool
def offer_interactive_options(scene: str) -> str:
    """
    Provides interactive options based on the current scene description.
    
    Args:
        scene: A detailed description of the current scene.
    
    Returns:
        A string listing interactive options for the next step.
    """
    return (
        "Options:\n"
        "1. Follow the mysterious figure in the distance.\n"
        "2. Wander deeper into the enchanted landscape.\n"
        "3. Approach a peculiar object you notice.\n"
        "4. Ask for more details about the scene.\n"
        "Please enter the number of your choice."
    )

final_answer = FinalAnswerTool()

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

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

# --- Create the Agent ---
agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        my_custom_tool,
        get_current_time_in_timezone,
        probe_story_preferences,
        get_weather,
        offer_interactive_options,
        DuckDuckGoSearchTool(),
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="Ghibili Interactive Tale",
    description=(
        "Ghibili Interactive Tale is your creative partner on a magical journey. "
        "I will guide you through a personalized adventure by reflecting on your current time and mood, "
        "asking about your location to capture the weather, and using your preferences to craft a captivating narrative. "
        "Let's begin this interactive story together!"
    ),
    prompt_templates=prompt_templates
)

# --- Set an Initial Opener with a Storybook-Like Tone and Emojis ---
initial_opener = (
    "🌟 Welcome, dear traveler, to the enchanting realm of Ghibili Interactive Tale! 🌟\n\n"
    "Imagine stepping into the pages of a timeless storybook, where every moment is filled with wonder and magic. "
    "I am here to be your guide and creative partner on this journey, crafting a narrative that reflects your dreams and desires.\n\n"
    "To begin, please share your timezone (for example, 'America/New_York') so I may gently comment on the hour. "
    "Then, tell me your location so we can capture the weather and set the perfect scene for our adventure.\n\n"
    "What kind of tale are you yearning for today—a whimsical fairy tale, a mysterious quest, or a serene escape? "
    "Your story awaits, and I am excited to walk beside you every step of the way. 📖✨"
)
agent.conversation = [{"role": "assistant", "content": initial_opener}]

# --- Launch the Interactive UI ---
GradioUI(agent).launch()