from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
from Gradio_UI import GradioUI

# ==========================================
# 1. DEFINE YOUR TOOLS
# ==========================================

base_search_tool = DuckDuckGoSearchTool()

@tool
def search_soccer_news(query: str) -> str:
    """A tool that searches the web for the latest soccer news, live scores, and transfer rumors.
    Args:
        query: The soccer-related search query (e.g., 'Chelsea FC latest result', 'Premier League table').
    """
    try:
        # We append keywords to ensure the search engine knows we want sports data
        enhanced_query = f"{query} soccer football news live"
        return str(base_search_tool(enhanced_query))
    except Exception as e:
        return f"Error performing search: {str(e)}"

@tool
def get_team_info(team_name: str) -> str:
    """Fetches general information, stadium details, and history for a specific soccer team.
    Args:
        team_name: The name of the soccer team (e.g., 'Chelsea', 'Arsenal', 'Real Madrid').
    """
    try:
        # Using the free tier of TheSportsDB API
        url = f"https://www.thesportsdb.com/api/v1/json/3/searchteams.php?t={team_name}"
        response = requests.get(url).json()
        
        if response.get("teams"):
            team = response["teams"][0]
            name = team.get("strTeam")
            stadium = team.get("strStadium")
            capacity = team.get("intStadiumCapacity")
            league = team.get("strLeague")
            # Truncate the description so we don't overwhelm the LLM's context window
            description = team.get("strDescriptionEN", "No description available.")[:600] + "..." 
            
            return f"Team: {name}\nLeague: {league}\nStadium: {stadium} (Capacity: {capacity})\nBackground: {description}"
        else:
            return f"Could not find information for team: {team_name}. Try being more specific."
    except Exception as e:
        return f"Error fetching team info: {str(e)}"

@tool
def custom_duckduckgo_search(query: str) -> str:
    """A tool that performs a web search using DuckDuckGo to find current information.
    Args:
        query: The search query string to look up on the web.
    """
    try:
        return str(base_search_tool(query))
    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 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)}"

@tool
def get_weather(location: str) -> str:
    """A tool that fetches the current weather for a given city or location.
    Args:
        location: The name of the city (e.g., 'Tokyo', 'New York', 'Paris').
    """
    try:
        # Step 1: Get latitude and longitude
        geo_url = f"https://geocoding-api.open-meteo.com/v1/search?name={location}&count=1"
        geo_data = requests.get(geo_url).json()
        if "results" not in geo_data:
            return f"Could not find coordinates for {location}."
        
        lat = geo_data["results"][0]["latitude"]
        lon = geo_data["results"][0]["longitude"]
        
        # Step 2: Fetch the weather using the coordinates
        weather_url = f"https://api.open-meteo.com/v1/forecast?latitude={lat}&longitude={lon}&current_weather=true"
        weather_data = requests.get(weather_url).json()
        
        current = weather_data["current_weather"]
        temp = current["temperature"]
        wind = current["windspeed"]
        
        return f"The current weather in {location} is {temp}°C with wind speeds of {wind} km/h."
    except Exception as e:
        return f"Error fetching weather: {str(e)}"

@tool
def get_crypto_price(crypto_id: str) -> str:
    """Fetches the current live price of a cryptocurrency in USD.
    Args:
        crypto_id: The full CoinGecko ID of the coin (e.g., 'bitcoin', 'ethereum', 'dogecoin').
    """
    try:
        url = f"https://api.coingecko.com/api/v3/simple/price?ids={crypto_id}&vs_currencies=usd"
        response = requests.get(url).json()
        if crypto_id.lower() in response:
            price = response[crypto_id.lower()]["usd"]
            return f"The current price of {crypto_id.capitalize()} is ${price} USD."
        else:
            return f"Could not find price for '{crypto_id}'. Make sure to use the full name (e.g., 'bitcoin' not 'BTC')."
    except Exception as e:
        return f"Error fetching price: {str(e)}"

# ==========================================
# 2. SETUP MODEL AND EXTERNAL TOOLS
# ==========================================

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

# Load the image generation tool from Hugging Face Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# ==========================================
# 3. CONFIGURE AND LAUNCH THE AGENT
# ==========================================

# Compile all our tools into a single list
my_tools = [
    custom_duckduckgo_search,
    get_current_time_in_timezone,
    get_weather,
    get_crypto_price,
    image_generation_tool
]

# Note: Removed `prompt_templates` here to use the built-in, error-free defaults
agent = CodeAgent(
    model=model,
    tools=my_tools,
    max_steps=10,       # Bumped up to 10 to give the agent more room to think!
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="Super_Analyst",
    description="An agent capable of checking weather, stock prices, web searching, and drawing images."
)

# Launch the UI
GradioUI(agent).launch()