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

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_cutom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"

# Tool to fetch latest news headlines from different news outlets
@tool
def search_news_headlines(source: str) -> str:
    """Searches for recent news headlines from specified source using DuckDuckGo.
    
    Args:
        source: News source from supported list (BBC News, CNN, Reuters, AP, 
               Fox News, CBS News, Wall Street Journal, Daily Mail)
    
    Returns:
        Formatted news results or error message
    """
    domain_map = {
        "BBC News": "bbc.com/news",
        "CNN": "cnn.com",
        "Reuters": "reuters.com",
        "AP": "apnews.com",
        "Fox News": "foxnews.com",
        "CBS News": "cbsnews.com",
        "Wall Street Journal": "wsj.com",
        "Daily Mail": "dailymail.co.uk"
    }
    
    try:
        domain = domain_map.get(source)
        if not domain:
            return f"Error: {source} is not a supported news source"
        
        # Simplified search query
        search_query = f"site:{domain} when:7d"
        
        # Use existing DuckDuckGoSearchTool implementation
        ddg_tool = DuckDuckGoSearchTool(max_results=5)
        raw_results = ddg_tool.forward(search_query)
        
        # Parse the raw results to extract titles and links
        import re
        headlines = re.findall(r'\[(.*?)\]\((.*?)\)', raw_results)
        
        # Filter out navigation/section pages but with fewer restrictions
        filtered_headlines = []
        for title, link in headlines:
            if not any(x in title.lower() for x in ['video clips', 'contact us', 'about us']):
                filtered_headlines.append((title, link))
        
        # Format the results
        formatted_headlines = [f"{i+1}. {title} {link}" for i, (title, link) in enumerate(filtered_headlines[:5])]
        
        if not formatted_headlines:
            return f"No recent news articles found from {source}. Please try again."
            
        return f"Here are the latest news from {source}:\n\n" + "\n".join(formatted_headlines)
    
    except Exception as e:
        return f"News search failed: {str(e)}"
        

final_answer = FinalAnswerTool()
model = HfApiModel(
max_tokens=2096,
temperature=0,
model_id='https://jc26mwg228mkj8dw.us-east-1.aws.endpoints.huggingface.cloud',# it is possible that this model may be overloaded
custom_role_conversions=None,
)

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

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, search_news_headlines], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()