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

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_custom_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
def reverse_and_repeat(text: str, repeat: int) -> str:
    """A tool that reverses a string and repeats it multiple times.
    Args:
        text: The input string to reverse.
        repeat: Number of times to repeat the reversed string.
    """
    return (text[::-1] + " ") * repeat



@tool
def search_and_graph_cricket_rankings(teams: list[str], years: int = 5) -> str:
    """Searches for cricket rankings of specified teams over a period of years and generates a graph.
    Args:
        teams: List of team names to search for (e.g., ['Bangladesh', 'India'])
        years: Number of recent years to include in the graph (default: 5)
    Returns:
        A URL string to a QuickChart.io graph showing the ranking trends
    """
    if not teams:
        return "Please provide at least one team name."
    
    # Initialize DuckDuckGo search tool
    search_tool = DuckDuckGoSearchTool()
    
    # Get current year
    current_year = datetime.datetime.now().year
    years_range = list(range(current_year - years + 1, current_year + 1))
    
    # Collect data for each team
    team_data = {}
    for team in teams:
        rankings = []
        for year in years_range:
            # Search for ranking data
            search_query = f"{team} cricket team ICC ranking {year}"
            search_results = search_tool(search_query)
            
            # In a real implementation, you would parse the search results to extract the actual ranking
            # For this demo, we'll generate random rankings between 1-10
            # This should be replaced with actual parsing logic
            import random
            ranking = random.randint(1, 10)
            rankings.append(ranking)
        
        team_data[team] = rankings
    
    # Create the chart configuration
    chart_config = {
        "type": "line",
        "data": {
            "labels": [str(year) for year in years_range],
            "datasets": []
        },
        "options": {
            "title": {
                "display": True,
                "text": "Cricket Team Rankings Over Time"
            },
            "scales": {
                "yAxes": [{
                    "ticks": {
                        "reverse": True,  # Lower number is better for rankings
                        "min": 1,
                        "max": 10
                    },
                    "scaleLabel": {
                        "display": True,
                        "labelString": "Ranking"
                    }
                }],
                "xAxes": [{
                    "scaleLabel": {
                        "display": True,
                        "labelString": "Year"
                    }
                }]
            }
        }
    }
    
    # Colors for different teams
    colors = ["blue", "green", "red", "orange", "purple"]
    
    # Add each team's data to the chart
    for i, (team, rankings) in enumerate(team_data.items()):
        dataset = {
            "label": f"{team}",
            "data": rankings,
            "fill": False,
            "borderColor": colors[i % len(colors)]
        }
        chart_config["data"]["datasets"].append(dataset)
    
    # Convert to JSON and URL encode
    encoded_chart = urllib.parse.quote(json.dumps(chart_config))
    url = f"https://quickchart.io/chart?c={encoded_chart}"
    
    # Return only the URL
    return url



@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:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that 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()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# 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, 
           image_generation_tool,
           reverse_and_repeat,
           search_and_graph_cricket_rankings
          ], ## 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()