from smolagents import CodeAgent, HfApiModel, load_tool, tool
import datetime
import pytz
import yaml
from transformers import pipeline
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Tool: Convert time between time zones
@tool
def convert_time(time_str: str, from_tz: str, to_tz: str) -> str:
    """Convert time from one timezone to another.
    Args:
        time_str: Time in 'YYYY-MM-DD HH:MM:SS' format.
        from_tz: Original timezone.
        to_tz: Target timezone.
    """
    try:
        from_zone = pytz.timezone(from_tz)
        to_zone = pytz.timezone(to_tz)
        local_time = datetime.datetime.strptime(time_str, "%Y-%m-%d %H:%M:%S")
        local_time = from_zone.localize(local_time)
        converted_time = local_time.astimezone(to_zone).strftime("%Y-%m-%d %H:%M:%S")
        return f"Time in {to_tz}: {converted_time}"
    except Exception as e:
        return f"Error converting time: {str(e)}"


# Tool: Assign alerts to departments based on keywords
@tool
def assign_alert(alert_message: str) -> str:
    """Assign an alert to a department based on keywords.
    Args:
        alert_message: The alert description.
    """
    alert_keywords = {
        "network": "Network Operations",
        "server": "Infrastructure",
        "database": "Database Team",
        "security": "Security Operations",
        "application": "Development Team"
    }
    for keyword, department in alert_keywords.items():
        if keyword.lower() in alert_message.lower():
            return f"Alert assigned to: {department}"
    return "Alert could not be assigned automatically. Please check manually."

final_answer = FinalAnswerTool()


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

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[final_answer, convert_time, assign_alert],
    max_steps=6,
    verbosity_level=1,
    planning_interval=None,
    name="Monitoring Assistant",
    description="An agent that helps with alert assignment, time conversion, monitoring tasks.",
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()