from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import pandas as pd
import json

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


# -----------------------------
# TIME TOOL
# -----------------------------
@tool
def get_time(timezone: str) -> str:
    """
    Get current time in a timezone.

    Args:
        timezone: Timezone name like Asia/Kathmandu

    Returns:
        Current time string
    """
    try:
        tz = pytz.timezone(timezone)
        return datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
    except Exception as e:
        return str(e)


# -----------------------------
# CALCULATOR TOOL
# -----------------------------
@tool
def calc(expression: str) -> str:
    """
    Evaluate math expression.

    Args:
        expression: Math expression like 2+2*5

    Returns:
        Result of calculation
    """
    try:
        return str(eval(expression))
    except:
        return "error"


# -----------------------------
# CSV TOOL
# -----------------------------
@tool
def read_csv(url: str) -> str:
    """
    Read CSV file from URL.

    Args:
        url: CSV file URL

    Returns:
        CSV content
    """
    try:
        df = pd.read_csv(url)
        return df.to_string()
    except Exception as e:
        return str(e)


# -----------------------------
# JSON TOOL
# -----------------------------
@tool
def read_json(url: str) -> str:
    """
    Read JSON file from URL.

    Args:
        url: JSON file URL

    Returns:
        JSON content
    """
    try:
        r = requests.get(url)
        return json.dumps(r.json(), indent=2)[:3000]
    except Exception as e:
        return str(e)


final_answer = FinalAnswerTool()


# If model overloads use endpoint:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'

model = HfApiModel(
    max_tokens=2096,
    temperature=0.1,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)


# optional image tool
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,
        DuckDuckGoSearchTool(),
        get_time,
        calc,
        read_csv,
        read_json
    ],
    max_steps=12,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()