import gradio as gr
import openai


def generate_plantuml2(api_key, text):
    openai.api_key = api_key
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are ChatGPT, a large language model trained by OpenAI. Generate PlantUML code for the following use case or code in natural language.",
            },
            {"role": "user", "content": text},
        ],
    )
    print(response)
    return response["choices"][0]["message"]["content"]

def generate_plantuml(api_key, text):
    openai.api_key = api_key
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are ChatGPT, a large language model trained by OpenAI. Generate PlantUML code for an architecture that uses Azure services and icons for the following use case or code in natural language. Choose the actors appropriately as per the use case. Do not use '!define SPRITESURL https://raw.githubusercontent.com/rabelenda/cicon-plantuml-sprites/v1.1.0/sprites' as it is outdated.",
            },
            {"role": "user", "content": text},
        ],
    )
    print(response)
    return response["choices"][0]["message"]["content"]

sample_text = '''
!define AzurePuml https://raw.githubusercontent.com/plantuml-stdlib/Azure-PlantUML/master/dist
!includeurl AzurePuml/AzureCommon.puml

actor Customer
actor Restaurant

Customer -> AzureAPIManagement : Login
AzureAPIManagement -> AzureActiveDirectory : Authenticate User
AzureActiveDirectory -> AzureAPIManagement : Return User Info

Customer -> AzureAPIManagement : Place Order
AzureAPIManagement -> AzureFunctionApp : Process Order
AzureFunctionApp -> AzureCosmosDB : Store Order Data
AzureFunctionApp -> Restaurant : Send Order Details
Restaurant -> AzureFunctionApp : Update Order Status
AzureFunctionApp -> AzureCosmosDB : Update Order Data
AzureFunctionApp -> Customer : Send Order Status

AzureFunctionApp -> AzureNotificationHubs : Send Push Notification

legend right
  Online Food Ordering App Architecture
endlegend
'''

iface = gr.Interface(
    fn=generate_plantuml,
    inputs=[
        gr.inputs.Textbox(label="OpenAI API Key"),
        gr.inputs.Textbox(label="Enter use case or code in natural language")
    ],
    outputs=gr.outputs.Textbox(label="Generated PlantUML Code"),
    title="PlantUML Code Generator",
    description="Generate PlantUML code using OpenAI's GPT-3.5-Turbo",
)

iface.launch()