from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
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_zodiac_info(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 returns your Zodiac sign and Chinese zodiac animal based on birth month and year.
    
    Args:
        arg1: the first argument
        arg2: the second argument

    """
    try:
        # Validate month input
        if not isinstance(arg1, str):
            return "Error: Birth month must be a string (e.g., 'February')."
            
        # Convert month to lowercase for case-insensitive comparison
        arg1 = arg1.lower()
        
        # Define months and their numerical values for easier date comparison
        month_to_num = {
            "january": 1, "february": 2, "march": 3, "april": 4,
            "may": 5, "june": 6, "july": 7, "august": 8,
            "september": 9, "october": 10, "november": 11, "december": 12
        }
        
        # Get month number
        month_num = month_to_num.get(arg1)
        if month_num is None:
            return "Error: Invalid month name provided. Please use full month names (e.g., 'February')."
        
        # Validate year input
        if not isinstance(arg2, int):
            try:
                arg2 = int(arg2)
            except (ValueError, TypeError):
                return "Error: Birth year must be a valid integer."
                
        if arg2 < 1900 or arg2 > 2100:
            return "Error: Birth year should be between 1900 and 2100."
        
        # Determine Western zodiac sign
        if (month_num == 1 and 20 <= 31) or (month_num == 2 and 1 <= 18):
            zodiac_sign = "Aquarius"
        elif (month_num == 2 and 19 <= 29) or (month_num == 3 and 1 <= 20):
            zodiac_sign = "Pisces"
        elif (month_num == 3 and 21 <= 31) or (month_num == 4 and 1 <= 19):
            zodiac_sign = "Aries"
        elif (month_num == 4 and 20 <= 30) or (month_num == 5 and 1 <= 20):
            zodiac_sign = "Taurus"
        elif (month_num == 5 and 21 <= 31) or (month_num == 6 and 1 <= 20):
            zodiac_sign = "Gemini"
        elif (month_num == 6 and 21 <= 30) or (month_num == 7 and 1 <= 22):
            zodiac_sign = "Cancer"
        elif (month_num == 7 and 23 <= 31) or (month_num == 8 and 1 <= 22):
            zodiac_sign = "Leo"
        elif (month_num == 8 and 23 <= 31) or (month_num == 9 and 1 <= 22):
            zodiac_sign = "Virgo"
        elif (month_num == 9 and 23 <= 30) or (month_num == 10 and 1 <= 22):
            zodiac_sign = "Libra"
        elif (month_num == 10 and 23 <= 31) or (month_num == 11 and 1 <= 21):
            zodiac_sign = "Scorpio"
        elif (month_num == 11 and 22 <= 30) or (month_num == 12 and 1 <= 21):
            zodiac_sign = "Sagittarius"
        else:
            zodiac_sign = "Capricorn"
        
        # Determine Chinese zodiac animal
        # Chinese zodiac repeats every 12 years
        # Rat, Ox, Tiger, Rabbit, Dragon, Snake, Horse, Goat, Monkey, Rooster, Dog, Pig
        chinese_zodiac = ["Monkey", "Rooster", "Dog", "Pig", "Rat", "Ox", 
                          "Tiger", "Rabbit", "Dragon", "Snake", "Horse", "Goat"]
        
        # Calculate Chinese zodiac animal
        zodiac_animal = chinese_zodiac[arg2 % 12]
        
        # Return formatted string
        return f"You're a {zodiac_sign} and you were born in the year of the {zodiac_animal}."
    
    except Exception as e:
        return f"Error: An unexpected error occurred: {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:
        # 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_zodiac_info, 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()