from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
import os
import matplotlib.pyplot as plt
import networkx as nx
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def text_to_graph_tool(relationships: str) -> str:
    """Generates a directed network graph image from relationships.

    Args:
        relationships: Semicolon-separated edges like 'User -> WebServer: Requests; WebServer -> Database: Queries'
    """
    try:
        plt.figure(figsize=(9, 6))
        G = nx.DiGraph()

        lines = [r.strip() for r in relationships.replace("\n", ";").split(";") if r.strip()]
        edge_labels = {}

        for entry in lines:
            if "->" in entry:
                parts = entry.split("->")
                source = parts[0].strip()
                remainder = parts[1].strip()

                if ":" in remainder:
                    target, label = remainder.split(":", 1)
                    target, label = target.strip(), label.strip()
                else:
                    target, label = remainder, ""

                G.add_edge(source, target)
                if label:
                    edge_labels[(source, target)] = label

        if len(G.nodes) == 0:
            return "No valid relationships found. Use format: 'NodeA -> NodeB: Label'"

        pos = nx.spring_layout(G, k=1.0, seed=42)
        nx.draw_networkx_nodes(G, pos, node_size=2800, node_color="#4F46E5", alpha=0.9)
        nx.draw_networkx_edges(G, pos, arrowstyle="->", arrowsize=18, edge_color="#64748B", width=2.0)
        nx.draw_networkx_labels(G, pos, font_size=9, font_color="white", font_weight="bold")

        if edge_labels:
            nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=8, font_color="#DC2626")

        plt.axis("off")
        plt.tight_layout()

        os.makedirs("generated_graphs", exist_ok=True)
        file_path = f"generated_graphs/graph_{int(datetime.datetime.now().timestamp())}.png"
        plt.savefig(file_path, dpi=200, bbox_inches="tight")
        plt.close()

        return f"Graph successfully created and saved at: {file_path}"
    except Exception as e:
        return f"Error generating graph: {str(e)}"

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


final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
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,
        text_to_graph_tool,
        image_generation_tool,
        get_current_time_in_timezone
    ],
)


GradioUI(agent).launch()