import gradio as gr
import openai
import plantuml
import os

# Function to call GPT-3.5-turbo model
def generate_plantuml_code(api_key, prompt):
    openai.api_key = api_key
    response = openai.Completion.create(
        engine="gpt-3.5-turbo",
        prompt=prompt,
        max_tokens=150,
        n=1,
        stop=None,
        temperature=0.7,
    )
    plantuml_code = response.choices[0].text.strip()
    return plantuml_code

# Function to convert PlantUML code to SVG image
def plantuml_code_to_svg(plantuml_code):
    with open("diagram.txt", "w") as f:
        f.write(plantuml_code)

    plantuml.exec_cmd(["-tsvg", "diagram.txt"])

    with open("diagram.svg", "r") as f:
        svg = f.read()

    return svg

# Function to process user input and generate diagram
def create_diagram(api_key, description):
    plantuml_code = generate_plantuml_code(api_key, description)
    svg_image = plantuml_code_to_svg(plantuml_code)
    return plantuml_code, svg_image

# Gradio interface
api_key_input = gr.Textbox(label="Enter your OpenAI API key")
description_input = gr.Textbox(lines=5, label="Describe the diagram in natural language")
plantuml_code_output = gr.Textbox(lines=10, label="Generated PlantUML code")
diagram_output = gr.Image(plot=True, label="Diagram")

iface = gr.Interface(
    fn=create_diagram,
    inputs=[api_key_input, description_input],
    outputs=[plantuml_code_output, diagram_output],
)
iface.launch()