"""
Custom Tools for GAIA Agent
Includes web search, weather info, and Hugging Face Hub statistics.
"""

from smolagents import Tool, DuckDuckGoSearchTool
from huggingface_hub import list_models
import random

# Export tools
__all__ = [
    'DuckDuckGoSearchTool',
    'WeatherInfoTool',
    'HubStatsTool',
    'search_tool',
    'weather_info_tool',
    'hub_stats_tool'
]

# Initialize the DuckDuckGo search tool
search_tool = DuckDuckGoSearchTool()


class WeatherInfoTool(Tool):
    name = "weather_info"
    description = "Fetches weather information for a given location. Useful for questions about weather conditions."
    inputs = {
        "location": {
            "type": "string",
            "description": "The location to get weather information for (e.g., 'Paris', 'New York', 'London')."
        }
    }
    output_type = "string"

    def forward(self, location: str) -> str:
        """
        Get weather information for a location.
        
        Args:
            location: City or location name
            
        Returns:
            Weather information string
        """
        # Note: This is a simplified implementation
        # In production, you would integrate with a real weather API
        weather_conditions = [
            {"condition": "Sunny", "temp_c": 22, "humidity": 60},
            {"condition": "Cloudy", "temp_c": 18, "humidity": 70},
            {"condition": "Rainy", "temp_c": 15, "humidity": 85},
            {"condition": "Clear", "temp_c": 25, "humidity": 55},
            {"condition": "Windy", "temp_c": 20, "humidity": 65}
        ]
        
        data = random.choice(weather_conditions)
        return (
            f"Weather in {location}:\n"
            f"Condition: {data['condition']}\n"
            f"Temperature: {data['temp_c']}°C\n"
            f"Humidity: {data['humidity']}%"
        )


# Initialize the weather tool
weather_info_tool = WeatherInfoTool()


class HubStatsTool(Tool):
    name = "hub_stats"
    description = "Fetches model statistics from Hugging Face Hub. Useful for questions about AI models and their popularity."
    inputs = {
        "author": {
            "type": "string",
            "description": "The username or organization name on Hugging Face Hub (e.g., 'meta-llama', 'Qwen', 'mistralai')."
        }
    }
    output_type = "string"

    def forward(self, author: str) -> str:
        """
        Get the most popular model from a Hugging Face author.
        
        Args:
            author: Hugging Face username or organization
            
        Returns:
            Information about the most downloaded model
        """
        try:
            # List models from the specified author, sorted by downloads
            models = list(list_models(
                author=author,
                sort="downloads",
                direction=-1,
                limit=5
            ))
            
            if models:
                result = f"Top models by {author}:\n\n"
                for i, model in enumerate(models[:5], 1):
                    result += (
                        f"{i}. {model.id}\n"
                        f"   Downloads: {model.downloads:,}\n"
                        f"   Likes: {model.likes}\n\n"
                    )
                return result.strip()
            else:
                return f"No models found for author/organization '{author}'."
        except Exception as e:
            return f"Error fetching models for {author}: {str(e)}"


# Initialize the Hub stats tool
hub_stats_tool = HubStatsTool()