from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import os
import sys
from tools.final_answer import FinalAnswerTool
from tools.web_search import DuckDuckGoSearchTool
from tools.visit_webpage import VisitWebpageTool

from Gradio_UI import GradioUI


# 1 - 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:
        # Create timezone object using pytz library
        tz = pytz.timezone(timezone)
        
        # Get current time in that timezone and format it as a readable string
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        
        # Return formatted response with the timezone and current time
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        # Handle any errors that might occur (invalid timezone, etc.)
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


# 2- Image generation
@tool
def generate_image_from_text(prompt: str) -> str:
    """A tool that generates an image based on a text description.
    Args:
        prompt: A detailed text description of the image you want to generate.
    """
    try:
        # Call the image generation tool loaded from Hugging Face Hub
        # The tool is loaded further down in the script before it's used here
        result = image_generation_tool(prompt)
        
        # Return success message with the result (which should contain image URL or path)
        return f"Image generated successfully: {result}"
    except Exception as e:
        # Handle any errors that occur during image generation
        return f"Error generating image: {str(e)}"


# 3 - WEB SEARCH TOOL - Using the existing DuckDuckGoSearchTool class
# Initialize the DuckDuckGo search tool from the tools directory
web_search_tool = DuckDuckGoSearchTool()

@tool
def search_web(query: str) -> str:
    """A tool that searches the web using DuckDuckGo for information.
    Args:
        query: The search query to find information on the web.
    """
    try:
        # Execute the search query using DuckDuckGo
        search_results = web_search_tool(query)
        
        # Format and return the search results
        return f"Search results for '{query}':\n\n{search_results}"
    except Exception as e:
        # Handle any errors that occur during the search
        return f"Error searching the web: {str(e)}"

# FINAL ANSWER TOOL (REQUIRED)
# Initialize the final answer tool - required for the agent to provide final answers
final_answer = FinalAnswerTool()


# VISIT WEBPAGE TOOL - From tools directory
# Initialize the visit webpage tool
visit_webpage = VisitWebpageTool()


# TOOL INITIALIZATION AND VERIFICATION
# Ensure all tools are properly loaded and print status
print("Initializing tools...")
try:
    # Import the image generation tool from Hugging Face Hub
    # This tool will be used by the generate_image_from_text function
    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
    print("✓ Image generation tool loaded successfully")
except Exception as e:
    print(f"✗ Failed to load image generation tool: {str(e)}")
    # Provide a fallback if the image generation tool fails to load
    def image_generation_tool(prompt):
        return f"[MOCK IMAGE] Generated from prompt: {prompt}"

print("✓ Time zone tool initialized")
print("✓ Web search tool initialized")
print("✓ Visit webpage tool initialized")
print("✓ Final answer tool initialized")


# MODEL CONFIGURATION
print("Configuring model...")
# If the agent does not answer, the model is overloaded
# Alternative endpoint: 'https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
model = HfApiModel(
    max_tokens=2096,           # Maximum number of tokens in the response
    temperature=0.5,           # Controls randomness: lower = more deterministic
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # Using Qwen 2.5 Coder model
    custom_role_conversions=None,
)
print("✓ Model configured successfully")


# LOAD PROMPT TEMPLATES
print("Loading prompt templates...")
# Create an upload directory if it doesn't exist
os.makedirs("uploads", exist_ok=True)

# Load prompt templates from YAML file for consistent agent responses
try:
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
    print("✓ Prompt templates loaded successfully")
except Exception as e:
    print(f"✗ Failed to load prompt templates: {str(e)}")
    # Provide default empty templates if loading fails
    prompt_templates = {}
    

# AGENT CONFIGURATION
print("Configuring agent...")
agent = CodeAgent(
    model=model,
    tools=[
        get_current_time_in_timezone,  # Tool 1: Time zone tool
        generate_image_from_text,      # Tool 2: Image generation tool
        search_web,                    # Tool 3: Web search tool
        visit_webpage,                 # Tool 4: Visit webpage tool (added from tools directory)
        final_answer                   # Required final answer tool
    ],
    max_steps=6,                       # Maximum number of reasoning steps
    verbosity_level=1,                 # Level of detail in agent's output
    grammar=None,                      # No specific grammar constraints
    planning_interval=None,            # No specific planning interval
    name=None,                         # No custom agent name
    description=None,                  # No custom agent description
    prompt_templates=prompt_templates  # Using loaded prompt templates
)
print("✓ Agent configured successfully")

# LAUNCH THE GRADIO UI
print("Launching Gradio UI...")
# Start the Gradio interface with our configured agent and file upload directory
GradioUI(agent, file_upload_folder="uploads").launch(debug=True)