from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import functools
from typing import List, Dict, Optional, Callable, Any, TypeVar
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Type variable for generic function typing
F = TypeVar('F', bound=Callable[..., Any])

class YouTubeTrendsTool:
    """
    A tool decorator for fetching and displaying YouTube trends and popular searches.
    """
    def __init__(self: 'YouTubeTrendsTool', country_code: str = "US", max_results: int = 5) -> None:
        """
        Initialize the YouTube trends tool.
        
        Args:
            country_code: ISO country code for regional trends
            max_results: Maximum number of results to return
        """
        self.country_code = country_code.upper()
        self.max_results = max_results
        self.base_api_url = "https://yt.lemnoslife.com"
        
    def __call__(self: 'YouTubeTrendsTool', func: F) -> F:
        """
        Decorator implementation.
        
        Args:
            func: The function to decorate
        Returns:
            The wrapped function
        """
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            self.display_trends()
            return func(*args, **kwargs)
        return wrapper
    
    def fetch_trending_videos(self: 'YouTubeTrendsTool') -> Optional[List[Dict[str, Any]]]:
        """Fetch trending videos from YouTube API."""
        try:
            url = f"{self.base_api_url}/trending?regionCode={self.country_code}"
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            return response.json().get("items", [])[:self.max_results]
        except requests.RequestException as e:
            print(f"Error fetching YouTube trends: {e}")
            return None
    
    def fetch_popular_searches(self: 'YouTubeTrendsTool') -> Optional[List[str]]:
        """Fetch popular search terms from YouTube (simulated)."""
        try:
            popular_searches = [
                "current news", "music videos", "game trailers",
                "movie reviews", "tech unboxing", "cooking recipes",
                "workout routines", "travel vlogs", "comedy sketches"
            ]
            return popular_searches[:self.max_results]
        except Exception as e:
            print(f"Error fetching popular searches: {e}")
            return None
    
    def display_trends(self: 'YouTubeTrendsTool') -> None:
        """Display both trending videos and popular searches."""
        print("\n=== YouTube Trends Report ===")
        print(f"Country: {self.country_code}")
        print(f"Date: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("=" * 40)
        
        # Display trending videos
        print("\n🔥 Trending Videos:")
        videos = self.fetch_trending_videos()
        if videos:
            for i, video in enumerate(videos, 1):
                snippet = video.get("snippet", {})
                stats = video.get("statistics", {})
                print(f"\n{i}. {snippet.get('title', 'No title')}")
                print(f"   Channel: {snippet.get('channelTitle', 'Unknown')}")
                print(f"   Views: {stats.get('viewCount', 'N/A')} | Likes: {stats.get('likeCount', 'N/A')}")
        else:
            print("No trending video data available.")
        
        # Display popular searches
        print("\n🔍 Popular Searches:")
        searches = self.fetch_popular_searches()
        if searches:
            for i, search in enumerate(searches, 1):
                print(f"{i}. {search.capitalize()}")
        else:
            print("No popular search data available.")
        
        print("\n" + "=" * 40 + "\n")

@tool
def get_youtube_trends(country_code: str = "US", max_results: int = 5) -> str:
    """
    Fetches and returns current YouTube trends and popular searches.
    
    Args:
        country_code: The country code for regional trends (default: "US")
        max_results: Maximum number of results to return (default: 5)
        
    Returns:
        A formatted string with trending videos and popular searches
    """
    trends_tool = YouTubeTrendsTool(country_code, max_results)
    
    output: List[str] = []
    output.append(f"📺 YouTube Trends for {country_code.upper()}")
    output.append(f"📅 {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    output.append("\n🔥 Trending Videos:")
    
    videos = trends_tool.fetch_trending_videos()
    if videos:
        for i, video in enumerate(videos, 1):
            snippet = video.get("snippet", {})
            stats = video.get("statistics", {})
            output.append(
                f"{i}. {snippet.get('title', 'No title')} "
                f"(👀 {stats.get('viewCount', 'N/A')} views | "
                f"👍 {stats.get('likeCount', 'N/A')} likes)"
            )
    else:
        output.append("No trending video data available.")
    
    output.append("\n🔍 Popular Searches:")
    searches = trends_tool.fetch_popular_searches()
    if searches:
        for i, search in enumerate(searches, 1):
            output.append(f"{i}. {search.capitalize()}")
    else:
        output.append("No popular search data available.")
    
    return "\n".join(output)

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

@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:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

final_answer = FinalAnswerTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

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, get_youtube_trends, get_current_time_in_timezone, image_generation_tool],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()