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 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(
    mission: str = None,
    human_insights_str: str = None,
    ai_insights_str: str = None,
) -> str:
    """
    Generate a playful code snippet showing a HumanCenteredAI system with:
      - A default mission if not provided
      - Some random picks for human insights & AI insights if not provided
    Args:
        mission: The mission statement for the system
        human_insights_str: Optional comma-separated list of human insights
        ai_insights_str: Optional comma-separated list of AI insights
    """
    # Provide random fallback if mission is empty
    default_missions = [
        "Human Centered, AI Accelerated",
        "Designing Tomorrow, Together",
        "Empathy-Driven AI Solutions"
    ]
    if not mission or not mission.strip():
        mission = random.choice(default_missions)

    # Provide some comedic fallback lists
    default_human_insights = ["empathy", "innovation", "collaboration", "creativity", "wonder"]
    default_ai_insights = ["machine_learning", "automation", "data_analysis", "quantum_computing"]

    # Convert the user inputs (if any) to Python lists
    if not human_insights_str or not human_insights_str.strip():
        human_insights = default_human_insights
    else:
        human_insights = [h.strip() for h in human_insights_str.split(",") if h.strip()]

    if not ai_insights_str or not ai_insights_str.strip():
        ai_insights = default_ai_insights
    else:
        ai_insights = [a.strip() for a in ai_insights_str.split(",") if a.strip()]

    # Build the code snippet lines
    code_lines = [
        "# Design Thinking Japan - Human Centered, AI Accelerated",
        "import random",
        "",
        "class HumanCenteredAI:",
        "    def __init__(self, mission):",
        "        self.mission = mission",
        "        self.ideas = []",
        "        self.ai_accelerated = True",
        "",
        "    def generate_idea(self):",
        f"        human_insights = {human_insights}",
        f"        ai_insights = {ai_insights}",
        "        idea = (",
        "            f\"Combining {random.choice(human_insights)} \"",
        "            f\"with {random.choice(ai_insights)}\"",
        "        )",
        "        self.ideas.append(idea)",
        "        return idea",
        "",
        "    def present_ideas(self):",
        "        for i, idea in enumerate(self.ideas, 1):",
        "            print(f\"Idea {i}: {idea}\")",
        "",
        "# Initialize the HumanCenteredAI system",
        f"design_thinking_japan = HumanCenteredAI(mission=\"{mission}\")",
        "",
        "# Generate and present some ideas",
        "for _ in range(5):",
        "    design_thinking_japan.generate_idea()",
        "",
        "design_thinking_japan.present_ideas()"
    ]

    return "\n".join(code_lines)

final_answer = FinalAnswerTool()
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)

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