import gradio as gr
import openai
import json
from graphviz import Digraph
import base64
from PIL import Image

def generate_knowledge_graph(api_key, user_input):
    print("Setting OpenAI API key...")
    openai.api_key = api_key

    print("Making API call to OpenAI...")
    completion = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-16k",
        messages=[
            {
                "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.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "metadata": {"type": "object"},
                        "nodes": {"type": "array"},
                        "edges": {"type": "array"}
                    },
                    "required": ["nodes", "edges"]
                }
            }
        ],
        function_call={"name": "knowledge_graph"}
    )

    print("Received response from OpenAI.")
    response_data = completion.choices[0]["message"]["function_call"]["arguments"]
    print(f"Response data: {response_data}")

    print("Converting response to JSON...")
    response_dict = json.loads(response_data)

    print("Generating knowledge graph using Graphviz...")
    dot = Digraph(comment="Knowledge Graph")

    # Add nodes to the graph
    for node in response_dict.get("nodes", []):
        dot.node(node["id"], f"{node['label']} ({node['type']})")

    # Add edges to the graph
    for edge in response_dict.get("edges", []):
        dot.edge(edge["from"], edge["to"], label=edge["relationship"])

    # Render to PNG format
    print("Rendering graph to PNG format...")
    dot.format = "png"
    dot.render(filename="knowledge_graph", cleanup=True)

    # Convert PNG to base64 to display in Gradio
    print("Converting PNG to base64...")
    with open("knowledge_graph.png", "rb") as img_file:
        img_base64 = base64.b64encode(img_file.read()).decode()

    print("Returning base64 image to Gradio interface.")
    return f"data:image/png;base64,{img_base64}"

iface = gr.Interface(
    fn=generate_knowledge_graph,
    inputs=[
        gr.inputs.Textbox(label="OpenAI API Key", type="password"),
        gr.inputs.Textbox(label="Text to Generate Knowledge Graph")
    ],
    outputs=gr.outputs.Image(type="pil", label="Generated Knowledge Graph"),
    live=False
)

print("Launching Gradio interface...")
iface.launch()