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

from Gradio_UI import GradioUI

rapid_api_key=os.environ["RAPID_KEY"]

# Below is an example of a tool that does nothing. Amaze us with your creativity !
#@tool
#def reflecting_on_results(arg1: str) -> str:
#    """A tool that receives the football teams comparison and then it reflects on the results with an insightful comment, it could be funny or witty
#    Args:
#        arg1: The first argument receives a string with the comparison of two teams as queried by user
#    """
#    insights
#    return insights



@tool
def get_football_results(team1: str, team2: str) -> str:
    """A tool that connects to an API that provides live football data and compares the stats of two teams in the English Premier League, then uses an LLM to generate a reflection on how both teams stack up.
    Args:
        team1: The first football team to check
        team2: The second football team to check
    """

    # First attempt to get actual data from the API
    url = "https://free-api-live-football-data.p.rapidapi.com/football-get-standing-home"
    querystring = {"leagueid": "47"}  # British Premier League
    headers = {
        "x-rapidapi-key": rapid_api_key,
        "x-rapidapi-host": "free-api-live-football-data.p.rapidapi.com"
    }

    # Try to get the data from the API
    try:
        response = requests.get(url, headers=headers, params=querystring)
        response_text = response.text  # Get the raw response as text

        # Load the TinyLlama model for analysis
        tiny_llama = pipeline("text-generation",
                              model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
                              torch_dtype=torch.bfloat16,
                              device_map="auto")

        # Create the prompt for the LLM
        messages = [
            {
                "role": "system",
                "content": f"You are a friendly AI system embedded in a wider AI agent. You are going to facilitate the work of other LLMs. You have received from user an API response in JSON format. In the response item 28 there is a long text starting with a text like 'status:success,response...You are going to extract the team names from the response and compare the stats of the two teams. First team check is {team1} and second team is {team2}. Provide a funny reflection on your analysis of how these two teams stack up against each other. You should provide a text response"
            },
            {
                "role": "user",
                "content": response_text
            }
        ]

        # Apply chat template and generate response
        prompt = tiny_llama.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        output = tiny_llama(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)

        # Extract the generated text
        comparison = output[0]['generated_text']

        # Return the comparison as is (it's already a string)
        return comparison

    except Exception as e:
        # Handle errors gracefully
        return f"""
## {team1} vs {team2} Premier League Comparison

I couldn't complete the analysis due to an error: {str(e)}

Let me provide some general information instead:

Both {team1} and {team2} are prominent teams in the English Premier League. To get a detailed comparison with current statistics, please try again later when the API and LLM processing are working properly.
"""

      

@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/QwQ-32B',# it is possible that this model may be overloaded, funny how it works with Qwen/QwQ-32B, default is Qwen/Qwen2.5-Coder-32B-Instruct
custom_role_conversions=None,
token=os.environ["HF_TOKEN"]
)


# 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, get_current_time_in_timezone, get_football_results], ## 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()