import gradio as gr
import openai
import json
from graphviz import Digraph

def generate_knowledge_graph(api_key, user_input):
    openai.api_key = api_key

    # Chamar a API da OpenAI
    print("Chamando a API da OpenAI...")
    completion = openai.Completion.create(
        model="gpt-3.5-turbo-16k",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Help me understand following by describing as a detailed knowledge graph: {user_input}",
            },
        ],
        functions=[
            {
                "name": "knowledge_graph",
                "description": "Generate a knowledge graph with entities and relationships. Use the colors to help differentiate between different node or edge types/categories. Always provide light pastel colors that work well with black font.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "metadata": {
                            "type": "object",
                            "properties": {
                                "createdDate": {"type": "string"},
                                "lastUpdated": {"type": "string"},
                                "description": {"type": "string"},
                            },
                        },
                        "nodes": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "id": {"type": "string"},
                                    "label": {"type": "string"},
                                    "type": {"type": "string"},
                                    "color": {"type": "string"},  # Added color property
                                    "properties": {
                                        "type": "object",
                                        "description": "Additional attributes for the node",
                                    },
                                },
                                "required": [
                                    "id",
                                    "label",
                                    "type",
                                    "color",
                                ],  # Added color to required
                            },
                        },
                        "edges": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "from": {"type": "string"},
                                    "to": {"type": "string"},
                                    "relationship": {"type": "string"},
                                    "direction": {"type": "string"},
                                    "color": {"type": "string"},  # Added color property
                                    "properties": {
                                        "type": "object",
                                        "description": "Additional attributes for the edge",
                                    },
                                },
                                "required": [
                                    "from",
                                    "to",
                                    "relationship",
                                    "color",
                                ],  # Added color to required
                            },
                        },
                    },
                    "required": ["nodes", "edges"],
                },
            }
        ],
        response_format="json",
    )
    response_data = completion.choices[0]["message"]["function_results"]["knowledge_graph"]
    print(response_data)

    # Visualizar o conhecimento usando Graphviz
    print("Gerando o conhecimento usando Graphviz...")
    dot = Digraph(comment="Knowledge Graph")
    for node in response_data.get("nodes", []):
        dot.node(node["id"], f"{node['label']} ({node['type']})", color=node.get("color", "lightblue"))
    for edge in response_data.get("edges", []):
        dot.edge(edge["from"], edge["to"], label=edge["relationship"], color=edge.get("color", "black"))

    # Renderizar para o formato PNG
    print("Renderizando o gráfico para o formato PNG...")
    dot.format = "png"
    dot.render(filename="knowledge_graph", cleanup=True)

    print("Gráfico gerado com sucesso!")

    return "knowledge_graph.png"

iface = gr.Interface(
    fn=generate_knowledge_graph,
    inputs=[
        gr.inputs.Textbox(label="OpenAI API Key", type="password"),
        gr.inputs.Textbox(label="User Input for Graph"),
    ],
    outputs=gr.outputs.Image(type="file", label="Generated Knowledge Graph"),
    live=False,
)

print("Iniciando a interface Gradio...")
iface.launch()