from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
import os
from tavily import TavilyClient
from Gradio_UI import GradioUI

import wikipedia
import textwrap

# Example tool template - can be modified for your needs

@tool
def my_custom_tool(arg1: str, arg2: int) -> str:  # Return type specification is important
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build?"


@tool
def wiki_search(query: str, sentences: int = 3) -> str:
    """Search Wikipedia and return a summary of the topic.
    Args:
        query: The topic to search for on Wikipedia
        sentences: Number of sentences to return in the summary (default: 3)
    Returns:
        str: A summary of the Wikipedia article
    """
    try:
        # Set language to English for consistency
        wikipedia.set_lang("en")
        
        # Search for the page
        search_results = wikipedia.search(query)
        if not search_results:
            return f"No Wikipedia results found for: {query}"
            
        try:
            # Get the page
            page = wikipedia.page(search_results[0], auto_suggest=False)
            
            # Get summary and full URL
            summary = wikipedia.summary(search_results[0], sentences=sentences, auto_suggest=False)
            url = page.url
            
            # Format the response
            response = f"Wikipedia article: {page.title}\n"
            response += f"URL: {url}\n\n"
            response += f"Summary:\n{textwrap.fill(summary, width=80)}\n"
            
            return response
            
        except wikipedia.DisambiguationError as e:
            # Handle disambiguation pages
            options = e.options[:5]  # Show first 5 options
            response = f"'{query}' may refer to multiple topics. Here are some options:\n\n"
            for i, option in enumerate(options, 1):
                response += f"{i}. {option}\n"
            return response
            
    except Exception as e:
        return f"Error searching Wikipedia: {str(e)}"

        
# @tool
# def tavily_search(query: str) -> str:
#     """A tool that performs web search using Tavily API and formats the results.
#     Args:
#         query: The search query to look up
#     Returns:
#         str: A formatted summary of search results
#     """
#     try:
#         client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
#         search_result = client.search(query=query, search_depth="advanced")
        
#         # Format the results in a user-friendly way
#         summary = f"Here are the search results for: {query}\n\n"
#         for result in search_result['results'][:5]:  # Top 5 results
#             summary += f"• {result['title']}\n"
#             summary += f"  Summary: {result['description'][:200]}...\n\n"
            
#         return summary
#     except Exception as e:
#         return f"Error performing search: {str(e)}"


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches and formats the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York', 'Asia/Tokyo')
    Returns:
        str: A formatted string containing the timezone and current time
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        # Return a more descriptive string with cleaned timezone name
        city = timezone.split('/')[-1].replace('_', ' ')
        return f"The current time in {city} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"
        

# Initialize tools
final_answer = FinalAnswerTool()
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# If the agent does not answer, the model is overloaded, please use another model or 
# the following Hugging Face Endpoint that also contains qwen2.5 coder:

# Model configuration
model_id = 'https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
model = HfApiModel(
    max_tokens=2096,
    temperature=0.3,  # Lower temperature for more consistent responses
    model_id=model_id,  # 'Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

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

# Initialize agent with tools
agent = CodeAgent(
    model=model,
    tools=[final_answer, wiki_search, get_current_time_in_timezone, image_generation_tool], # tavily_search,
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Launch the Gradio interface
GradioUI(agent).launch()