#!/usr/bin/env python
# coding=utf-8
"""
Robust startup file for your SmolAgents Space.
– Merges custom prompts.yaml with safe defaults
– Defines tools, model, CodeAgent, and Gradio UI
"""

import copy
import datetime
import pytz
import yaml

# --------- SmolAgents imports ---------
from smolagents import CodeAgent, HfApiModel, tool

# --------- Local Gradio wrapper --------
from Gradio_UI import GradioUI

# --------- Local tools -----------------
from tools.final_answer import FinalAnswerTool
from tools.visit_webpage import VisitWebpageTool
from tools.web_search import DuckDuckGoSearchTool

# ---------------------------------------------------------------------------
# Example custom tool (now with full doc-string so SmolAgents can JSON-schema it)
@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """Example Tool: a simple placeholder.

    Args:
        arg1 (str): Any string input (unused – demo only).
        arg2 (int): Any integer input (unused – demo only).

    Returns:
        str: A constant demo message.
    """
    return "What magic will you build ?"


# A small utility tool: get local time in a given timezone
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Return the current local time in the specified timezone.

    Args:
        timezone (str): IANA timezone name, e.g. 'America/New_York'.

    Returns:
        str: Formatted datetime string or an error message.
    """
    try:
        tz = pytz.timezone(timezone)
        return datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
    except Exception as exc:
        return f"Error: {exc}"


# ---------------------------------------------------------------------------
# Model configuration
model = HfApiModel(
    model_id="Qwen/Qwen3-32B",   # Adjust if overloaded
    max_tokens=2096,
    temperature=0.5,
    custom_role_conversions=None,
)

# ---------------------------------------------------------------------------
# 1) Load prompts.yaml (may be imperfect)
try:
    with open("prompts.yaml", "r") as f:
        custom_prompts = yaml.safe_load(f) or {}
except Exception as exc:
    print("⚠️  Could not load prompts.yaml – using bare defaults.", exc)
    custom_prompts = {}

# 2) Build minimal required template structure
MIN_SYSTEM = "You are a helpful agent."
MIN_FINAL = (
    "Thought: I now know the answer.\n"
    "Code:\n```py\nfinal_answer({{answer}})\n```<end_code>"
)
MIN_PLAN = {
    "initial_facts": "|-\n  List any facts.",
    "initial_plan": "|-\n  Write a high-level plan.",
    "update_facts_pre_messages": "|-\n  Pre facts.",
    "update_facts_post_messages": "|-\n  Post facts.",
    "update_plan_pre_messages": "|-\n  Pre plan.",
    "update_plan_post_messages": "|-\n  Post plan.",
}
MIN_MGR = {
    "task": "|-\n  You are {{name}}.",
    "report": "|-\n  {{final_answer}}",
}

def ensure(mapping, key, default):
    if key not in mapping or not isinstance(mapping[key], type(default)):
        mapping[key] = copy.deepcopy(default)

ensure(custom_prompts, "system_prompt", MIN_SYSTEM)
ensure(custom_prompts, "final_answer", MIN_FINAL)
ensure(custom_prompts, "planning",    MIN_PLAN)
ensure(custom_prompts, "managed_agent", MIN_MGR)

# If planning / managed_agent came in as strings, convert to dict default
if isinstance(custom_prompts["planning"], str):
    custom_prompts["planning"] = copy.deepcopy(MIN_PLAN)
if isinstance(custom_prompts["managed_agent"], str):
    custom_prompts["managed_agent"] = copy.deepcopy(MIN_MGR)

prompt_templates = custom_prompts  # now guaranteed correct structure

# ---------------------------------------------------------------------------
# Build the CodeAgent
agent = CodeAgent(
    model=model,
    tools=[
        FinalAnswerTool(),
        DuckDuckGoSearchTool(),
        VisitWebpageTool(),
        get_current_time_in_timezone,
        my_custom_tool,
    ],
    max_steps=6,
    verbosity_level=1,
    prompt_templates=prompt_templates,
)

# ---------------------------------------------------------------------------
# Launch the Gradio UI
if __name__ == "__main__":
    GradioUI(agent).launch()