import datetime
import pytz
import yaml
from smolagents import CodeAgent, HfApiModel
from smolagents.tools import tool
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
import datetime
import pytz
import yaml
import math
# ✅ This is the CORRECT and SIMPLE agent
from smolagents.agents import CodeAgent
from smolagents import HfApiModel
from smolagents.tools import tool

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
# -----------------------------
# ✅ Define tools directly here
# -----------------------------

@tool
def sigmoid_transform(x1: float, x2: float, x3: float) -> str:
    """Applies the sigmoid transformation to three numbers and returns markdown-formatted output.
    
    Args:
        x1: First number
        x2: Second number
        x3: Third number
    
    Returns:
        A markdown string with sigmoid-transformed values.
    """
    def sigmoid(x):
        return 1 / (1 + math.exp(-x))
    
    return (
        "### 🔥 Sigmoid Transformation Results\n"
        f"- **x1 = {x1}** → σ(x1) = `{sigmoid(x1):.4f}`\n"
        f"- **x2 = {x2}** → σ(x2) = `{sigmoid(x2):.4f}`\n"
        f"- **x3 = {x3}** → σ(x3) = `{sigmoid(x3):.4f}`\n"
    )

@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build?"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Fetch the current local time in a specified timezone.
    Args:
        timezone: A valid timezone string (e.g., 'America/New_York').
    """
    try:
        tz = pytz.timezone(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: {str(e)}"

# -----------------------------
# ✅ Load simplified prompt template
# -----------------------------
with open("prompts.yaml") as f:
    prompt_templates = yaml.safe_load(f)

# -----------------------------
# ✅ Set up model and agent
# -----------------------------
model = HfApiModel(
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    max_tokens=2048,
    temperature=0.5
)

final_answer = FinalAnswerTool()

agent = CodeAgent(
    model=model,
    tools=[final_answer, my_custom_tool, get_current_time_in_timezone],
    prompt_templates=prompt_templates,
    max_steps=6,
    verbosity_level=1,
)

# -----------------------------
# ✅ Launch Gradio interface
# -----------------------------
GradioUI(agent).launch()