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

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 ai_fortune_teller(category: str) -> str:
    """Predicts the future with extreme sarcasm and humor.
    
    Args:
        category: The category to predict ('career', 'love', 'finance').
    
    Returns:
        str: A sarcastic and funny prediction.
    """
    fortunes = {
        'career': [
            "Your dream job? Oh, you mean winning the lottery and never working again? Keep dreaming. 😂",
            "You're destined for greatness... if scrolling social media for 5 hours a day counts as a skill.",
            "One day, you will lead a company. That company may just be your own failed startup, but hey, leadership is leadership. 👏",
            "Your boss will finally appreciate you! Just kidding. They’ll steal your ideas and take all the credit. Classic. 🤡",
            "A new job opportunity is coming your way. Too bad it requires 15 years of experience for an entry-level position. 🤦‍♂️"
        ],
        'love': [
            "What is love? I was hoping you’d tell me. Go ask your mom, she might know. 😂",
            "You will meet the love of your life... and they will promptly ghost you after three messages. 💔",
            "A grand romantic gesture is in your future! Unfortunately, it’s from someone you find incredibly annoying. 🤷",
            "Your soulmate is out there... probably swiping left on your profile as we speak. 😬",
            "True love is coming. It might just be your delivery guy, but hey, a connection is a connection. 🚀"
        ],
        'finance': [
            "Your financial future is bright! Unless, of course, you continue buying overpriced coffee every day. ☕💸",
            "You will soon receive a large sum of money. Oh wait, I misread. You will soon receive a large sum of bills. Oops. 🤡",
            "Invest in crypto! Or don’t. Either way, you’ll still be confused about how it actually works. 🧐",
            "You are about to make a life-changing financial decision... by finally unsubscribing from that gym membership you never use. 👏",
            "A golden financial opportunity is coming. Unfortunately, it’s an MLM pitch from your cousin. 🚩"
        ]
    }

    if category in fortunes:
        return random.choice(fortunes[category])
    else:
        return "Oh, you want a special prediction? Sure. I predict you’ll ignore this advice and continue making questionable life choices. 😂"


@tool
def conspiracy_theory_generator() -> str:
    """Generates a random, ridiculous conspiracy theory.
    
    Returns:
        str: A completely fake but hilarious conspiracy theory.
    """
    subjects = ["AI", "Elon Musk", "NASA", "The Moon", "Time Travelers", "Ancient Egyptians", "Pigeons", "The Illuminati"]
    events = ["never actually happened", "was staged to hide something bigger", "is secretly controlled by lizard people", 
              "is part of a government experiment", "is an AI simulation that we’re all trapped in"]
    twists = ["and the government knows about it", "but it's being covered up", "and it’s all connected to the Bermuda Triangle", 
              "but nobody is ready for the truth yet", "and it will all be revealed in 2030"]

    return f"{random.choice(subjects)} {random.choice(events)} {random.choice(twists)}."
        


@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)

# initialize duckduck go search tool
duckduckgosearch_tool = DuckDuckGoSearchTool()

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, image_generation_tool, duckduckgosearch_tool, ai_fortune_teller, conspiracy_theory_generator], ## 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()