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 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 my_super_tool(city: str) -> str:
    """A tool that does something interesting with a city.
    Args:
        city: The name of a city
    """
    return f"I will do something special with {city}!"

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

@tool
def generate_futuristic_code_snippet(
    company_name: str,
    mission: str = None,
    key_values: str = None
) -> str:
    """
    Generates a futuristic-style Python code snippet given a company name,
    an optional mission, and an optional set of key values.

    Args:
        company_name: The name of the company (e.g. "Design Thinking Japan")
        mission: The company's mission statement (e.g. "Human Centered, AI Accelerated").
                 If not provided, a random futuristic mission is chosen.
        key_values: A comma-separated string of key values (e.g. "empathy, innovation").
                    If not provided, a set of futuristic key values is randomly picked.
    """
    # 1. Provide random defaults if mission/key_values are missing
    random_missions = [
        "Automate the intergalactic supply chain",
        "Enlighten minds through quantum synergy",
        "Transcend beyond the boundaries of reality",
        "Engineer post-human empathy"
    ]
    random_values = [
        "hyper-collaboration",
        "cosmic curiosity",
        "matrix-like agility",
        "singularity-driven passion"
    ]
    
    if mission is None or not mission.strip():
        mission = random.choice(random_missions)
    
    if key_values is None or not key_values.strip():
        # pick 2–3 random values
        chosen_values = random.sample(random_values, k=2)
        key_values = ", ".join(chosen_values)
    
    # 2. Prepare the snippet lines
    code_lines = [
        f"def {company_name.lower().replace(' ', '_')}_futuristic():",
        f"    \"\"\"",
        f"    This function encapsulates {company_name}'s futuristic ambitions.",
        f"    Mission: {mission}",
        f"    Key Values: {key_values}",
        f"    \"\"\"",
        "    # Initiating hyper-drive synergy with mission parameters",
        "    synergy_level = 9999",
        "    dimensional_rift = True",
        "",
        "    if dimensional_rift:",
        "        print(f'Breach opened! {company_name} is transcending conventional limits...')",
        "    else:",
        "        raise Exception('No rift detected. Expansion halted!')",
        "",
        f"    print(f'{company_name} pushing forward with unstoppable momentum...')",
        "",
        f"    return f'{company_name} has successfully infused the future with {mission}!'",
    ]
    
    # 3. Combine the lines into a single string
    snippet = "\n".join(code_lines)
    return snippet

final_answer = FinalAnswerTool()
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='https://wxknx1kg971u7k1n.us-east-1.aws.endpoints.huggingface.cloud',# 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,                # Must always keep this
        my_custom_tool,             # The dummy one from the snippet
        get_current_time_in_timezone,  # The time-zone tool
        my_super_tool,
        generate_futuristic_code_snippet], ## 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()