import gradio as gr
import openai
import json
from graphviz import Digraph
import base64
from io import BytesIO
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.Completion.create(
        engine="text-davinci-002",
        prompt=f"Help me understand the following by describing it as a detailed knowledge graph: {user_input}",
        max_tokens=100
    )

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

    # For demonstration, let's assume the response_data is a JSON string that can be converted to a dictionary.
    # You'll need to write code to interpret the text-based response to generate this dictionary.
    print("Converting response to JSON...")
    try:
        response_dict = json.loads(response_data)
    except json.JSONDecodeError:
        print("Failed to decode JSON. Using empty dictionary as a fallback.")
        response_dict = {}

    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()