from smolagents import CodeAgent, InferenceClientModel, load_tool, tool
import datetime
import pytz
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Get the current local time in the requested timezone.

    Args:
        timezone: IANA timezone name (e.g. 'Asia/Kolkata', 'America/New_York', 'Europe/London')

    Returns:
        Human-readable string with current time or error message
    """
    try:
        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz)
        return f"Current time in **{timezone}**: {now:%Y-%m-%d %H:%M:%S %Z}"
    except pytz.exceptions.UnknownTimeZoneError:
        return f"Unknown timezone: '{timezone}' — please use a valid IANA name"
    except Exception as e:
        return f"Error: {str(e)}"


@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """Placeholder custom tool — replace with real logic as needed.

    Args:
        arg1: some text input (e.g. a name, description, or keyword)
        arg2: some integer input (e.g. a count, year, or multiplier)

    Returns:
        Demonstration response string
    """
    return f"Tool called → arg1 = {arg1!r} | arg2 = {arg2} → What would you like to build?"


def main():
    final_answer_tool = FinalAnswerTool()

    model = InferenceClientModel(
        model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
        temperature=0.5,
        max_tokens=2048,
    )

    # Optional: text-to-image tool
    image_tool = None
    try:
        image_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
        print("Text-to-image tool loaded")
    except Exception as e:
        print(f"Text-to-image load failed: {e}")

    tools = [
        final_answer_tool,
        get_current_time_in_timezone,
        my_custom_tool,
    ]
    if image_tool is not None:
        tools.append(image_tool)

    agent = CodeAgent(
        tools=tools,
        model=model,
        max_steps=10,
        verbosity_level=1,
        stream_outputs=True,
        # planning_interval=4,           # optional
        # prompt_templates=...           # only if you really need custom prompts
    )

    GradioUI(agent).launch(
        server_name="0.0.0.0",
        server_port=7860,
        # share=True,
    )


if __name__ == "__main__":
    main()