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

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def get_pokemon_locations(pokemon_name: str) -> dict:
     
    """A tool that retrieves the areas where Pokémon can be found.
    Args:
        pokemon_name: the name of the pokemon
    Returns:
        A json with the locations, here an example of a json response:
        [
          {
            "location_area": {
              "name": "kanto-route-2-south-towards-viridian-city",
              "url": "https://pokeapi.co/api/v2/location-area/296/"
            },
            "version_details": [
              {
                "max_chance": 10,
                "encounter_details": [
                  {
                    "min_level": 7,
                    "max_level": 7,
                    "condition_values": [
                      {
                        "name": "time-morning",
                        "url": "https://pokeapi.co/api/v2/encounter-condition-value/3/"
                      }
                    ],
                    "chance": 5,
                    "method": {
                      "name": "walk",
                      "url": "https://pokeapi.co/api/v2/encounter-method/1/"
                    }
                  }
                ],
                "version": {
                  "name": "heartgold",
                  "url": "https://pokeapi.co/api/v2/version/15/"
                }
              }
            ]
          }
        ]
    """
    try:
        url = "https://pokeapi.co/api/v2/pokemon/"+pokemon_name+"/encounters"
        # Send a GET request to the URL
        response = requests.get(url)
        json_data = "{}"
        # Check if the request was successful
        if response.status_code == 200:
            # Try to parse the response as JSON
            try:
                json_data = response.json()
            except ValueError as e:
                print(f"Failed to parse JSON: {e}")
        else:
            print(f"Failed to retrieve data. Status code: {response.status_code}")
    except requests.exceptions.RequestException as e:
        print(f"An error occurred: {e}")

    return json_data

    
@tool
def get_pokemon_detail(pokemon_name:str)-> dict: #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 retrieves the characteristic of a pokemon
    Args:
        pokemon_name: the name of the pokemon
    Returns:
        A json with the pokemon details, here an example of a partial json response for the pokemon 'clefairy':
        {
          "id": 35,
          "name": "clefairy",
          "base_experience": 113,
          "height": 6,
          "is_default": true,
          "order": 56,
          "weight": 75,
          "abilities": [
            {
              "is_hidden": true,
              "slot": 3,
              "ability": {
                "name": "friend-guard",
                "url": "https://pokeapi.co/api/v2/ability/132/"
              }
            }
          ],
          "forms": [
            {
              "name": "clefairy",
              "url": "https://pokeapi.co/api/v2/pokemon-form/35/"
            }
          ],
          "game_indices": [
            {
              "game_index": 35,
              "version": {
                "name": "white-2",
                "url": "https://pokeapi.co/api/v2/version/22/"
              }
            }
          ],
          "held_items": [
            {
              "item": {
                "name": "moon-stone",
                "url": "https://pokeapi.co/api/v2/item/81/"
              },
              "version_details": [
                {
                  "rarity": 5,
                  "version": {
                    "name": "ruby",
                    "url": "https://pokeapi.co/api/v2/version/7/"
                  }
                }
              ]
            }
          ],
          "location_area_encounters": "/api/v2/pokemon/35/encounters",
          "moves": [
            {
              "move": {
                "name": "pound",
                "url": "https://pokeapi.co/api/v2/move/1/"
              },
              "version_group_details": [
                {
                  "level_learned_at": 1,
                  "version_group": {
                    "name": "red-blue",
                    "url": "https://pokeapi.co/api/v2/version-group/1/"
                  },
                  "move_learn_method": {
                    "name": "level-up",
                    "url": "https://pokeapi.co/api/v2/move-learn-method/1/"
                  }
                }
              ]
            }
          ],
          "species": {
            "name": "clefairy",
            "url": "https://pokeapi.co/api/v2/pokemon-species/35/"
          }
        }
    """
    try:
        url = "https://pokeapi.co/api/v2/pokemon/"+pokemon_name+"/"
        # Send a GET request to the URL
        response = requests.get(url)
        json_data = "{}"
        # Check if the request was successful
        if response.status_code == 200:
            # Try to parse the response as JSON
            try:
                json_data = response.json()
            except ValueError as e:
                print(f"Failed to parse JSON: {e}")
        else:
            print(f"Failed to retrieve data. Status code: {response.status_code}")
    except requests.exceptions.RequestException as e:
        print(f"An error occurred: {e}")

    return json_data

@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
    # model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
    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=[get_pokemon_locations, get_pokemon_detail, image_generation_tool, get_current_time_in_timezone, final_answer], ## 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()