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 AWS services and icons for the following use case or code in natural language.",
            },
            {"role": "user", "content": text},
        ],
    )
    print(response)
    return response["choices"][0]["message"]["content"]

sample_text = '''
!define AWSPUML https://raw.githubusercontent.com/awslabs/aws-icons-for-plantuml/v14.0/PlantUMLAWSDiagramGenerator.puml
!includeurl AWSPUML/AWSCommon.puml

actor Customer
actor Restaurant

Customer -> AWSAPIGateway : Login
AWSAPIGateway -> AWSIAM : Authenticate User
AWSIAM -> AWSAPIGateway : Return User Info

Customer -> AWSAPIGateway : Place Order
AWSAPIGateway -> AWSLambda : Process Order
AWSLambda -> AmazonDynamoDB : Store Order Data
AWSLambda -> Restaurant : Send Order Details
Restaurant -> AWSLambda : Update Order Status
AWSLambda -> AmazonDynamoDB : Update Order Data
AWSLambda -> Customer : Send Order Status

AWSLambda -> AmazonSNS : Send Push Notification

legend right
  Online Food Ordering App Architecture (AWS)
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()