from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool, Tool
import datetime
import requests
import pytz
import yaml
import gradio as gr

# ========== FINAL ANSWER TOOL ==========
class FinalAnswerTool(Tool):
    name = "final_answer"
    description = "Call this with the final answer to the user's question"
    inputs = {
        "answer": {
            "type": "string",
            "description": "The final answer to send back to the user"
        }
    }
    output_type = "string"
    
    def forward(self, answer: str):
        return answer
# =======================================

# Wikipedia Tool
@tool
def search_wikipedia(query: str, sentences: int = 3) -> str:
    """
    Search Wikipedia for information.
    
    Args:
        query (str): The topic to search for
        sentences (int, optional): Number of summary sentences to return. Defaults to 3.
    
    Returns:
        str: A summary from Wikipedia
    """
    try:
        url = "https://en.wikipedia.org/w/api.php"
        params = {
            "action": "query", "format": "json", "titles": query,
            "prop": "extracts", "exintro": True, "explaintext": True,
            "exsentences": sentences
        }
        response = requests.get(url, params=params, timeout=10)
        data = response.json()
        pages = data.get("query", {}).get("pages", {})
        for page_id, page_data in pages.items():
            if page_id != "-1":
                extract = page_data.get("extract", "No information found.")
                return f"**Wikipedia: {query}**\n\n{extract}"
        return f"No Wikipedia page found for '{query}'."
    except Exception as e:
        return f"Error: {str(e)}"

# Time Tool
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Get current local time in a specified timezone.
    
    Args:
        timezone (str): A valid timezone (e.g., 'America/New_York')
    
    Returns:
        str: The current local time
    """
    try:
        tz = pytz.timezone(timezone)
        time_str = datetime.datetime.now(tz).strftime("%Y-%m-%d %I:%M:%S %p")
        return f"⏰ Time in {timezone}: {time_str}"
    except:
        return f"Unknown timezone. Try: 'America/New_York', 'Europe/London', 'Asia/Tokyo'"

# Initialize
final_answer = FinalAnswerTool()

# Use smaller model to avoid overload
model = HfApiModel(
    model_id='Qwen/Qwen2.5-Coder-7B-Instruct',
    max_tokens=1024,
    temperature=0.3
)

# Simple tools list
tools_list = [
    final_answer, 
    DuckDuckGoSearchTool(), 
    search_wikipedia, 
    get_current_time_in_timezone
]

# Load prompts if exists
try:
    with open("prompts.yaml", 'r') as f:
        prompt_templates = yaml.safe_load(f)
except:
    prompt_templates = None

# Create agent
agent = CodeAgent(
    model=model,
    tools=tools_list,
    max_steps=6,
    name="Alfred",
    description="AI Assistant with Wikipedia search, web search, and time zone lookup"
)

# Gradio Interface
def chat_with_agent(message, history):
    """Process user message"""
    try:
        response = agent.run(message)
        return response
    except Exception as e:
        return f"Sorry, I encountered an error: {str(e)}"

# Create interface - REMOVED theme="soft"
demo = gr.ChatInterface(
    fn=chat_with_agent,
    title="🤖 Alfred AI Assistant",
    description="Ask me anything! I can search Wikipedia, search the web, and tell time in any timezone.",
    examples=[
        "Tell me about Albert Einstein",
        "What time is it in Tokyo, Japan?",
        "Search for latest AI news",
        "What is quantum computing?"
    ]
)

if __name__ == "__main__":
    demo.launch(share=False)