import gradio as gr
import openai
import json
from graphviz import Digraph
from PIL import Image
import io
import requests
from bs4 import BeautifulSoup
from ast import literal_eval

# Function to scrape text from a website
def scrape_text_from_url(url):
    response = requests.get(url)
    if response.status_code != 200:
        return "Error: Could not retrieve content from URL."
    soup = BeautifulSoup(response.text, "html.parser")
    paragraphs = soup.find_all("p")
    text = " ".join([p.get_text() for p in paragraphs])
    return text

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

    # Check if input is URL or text
    if user_input.startswith("http://") or user_input.startswith("https://"):
        user_input = scrape_text_from_url(user_input)

    # Chamar a API da 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. 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"],
                },
            }
        ],
        function_call={"name": "knowledge_graph"},
    )

    response_data = completion.choices[0]["message"]["function_call"]["arguments"]

    try:
        if isinstance(response_data, str):
            response_data = literal_eval(response_data)
    except (ValueError, SyntaxError) as e:
        print(f"Error in decoding JSON or literal_eval: {e}")
        return "Error in decoding JSON"

    if not isinstance(response_data, dict):
        print("Unexpected data type for response_data")
        return "Error: Unexpected data type"

    dot = Digraph(comment="Knowledge Graph", format='png')
    dot.attr(dpi='300')
    dot.attr(bgcolor='white')
    dot.attr('node', shape='box', style='filled', fillcolor='lightblue', fontcolor='black')

    for node in response_data.get("nodes", []):
        dot.node(node["id"], f"{node['label']} ({node['type']})", color=node.get("color", "lightblue"))

    dot.attr('edge', color='black', fontcolor='black')

    for edge in response_data.get("edges", []):
        dot.edge(edge["from"], edge["to"], label=edge["relationship"], color=edge.get("color", "black"))

    image_data = dot.pipe()
    image = Image.open(io.BytesIO(image_data))

    return image

title_and_description = """
# Instagraph - Knowledge Graph Generator
Created by [@artificialguybr](https://twitter.com/artificialguybr)  
Code by [Instagraph on GitHub](https://github.com/yoheinakajima/instagraph)

Enter your OpenAI API Key and a question, and let the AI create a detailed knowledge graph for you.

## Features
- **URL**: You can now input a URL to scrape text for generating the knowledge graph.
- **Security**: Rest assured, the code is open for your inspection in the files. There's no risk in using your OpenAI API key here.
- **Best View**: For the best visualization, consider downloading the generated image.
- **Flexible Input**: You can either type what you want the API to generate as a graph or use a URL for this purpose.

Feel free to explore and generate your own knowledge graphs!
"""


with gr.Blocks() as app:
    gr.Markdown(title_and_description)
    
    with gr.Row():
        with gr.Column():
            result_image = gr.Image(type="pil", label="Generated Knowledge Graph")
            
    with gr.Row():
        with gr.Column():
            api_key = gr.Textbox(label="OpenAI API Key", type="password")
            user_input = gr.Textbox(label="User Input for Graph or URL", type="text")
            run_btn = gr.Button("Generate")
            
    run_btn.click(
        generate_knowledge_graph,
        inputs=[api_key, user_input],
        outputs=[result_image]
    )

app.queue(concurrency_count=10)

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