import gradio as gr
import openai
import os
import requests

openai.organization = os.getenv("API_ORG")
openai.api_key = os.getenv("API_KEY")
app_password = os.getenv("APP_PASSWORD")
app_username = os.getenv("APP_USERNAME")

def generate_prompt(input):
    prompt = """You are a prompt writing support system for image generation AI.
From the user's input, You output prompt in English that should be input to the image generation AI, imagining its intent as much as possible.
You are not allowed to ask questions of the user.
You will always output only brief prompt in English to be input to the image generation AI.
Your output will always English.
Input from user:
"""
    response = openai.ChatCompletion.create(
        model = "gpt-3.5-turbo",
        messages = [{"role": "system", "content": prompt+input}],
        max_tokens=256
    )
    generated_text = response['choices'][0]['message']['content'].strip()
    return "Make the illustration a photo: "+generated_text

def get_related_caption(prompt):
    url = "https://api.irasutoya.nibo.sh/semantic-search"
    params = {'q': prompt}
    headers = {"content-type": "application/json"}
    r = requests.get(url, params=params, headers=headers)
    data = r.json()
    return data['illustrations'][0]['description']

def generate(prompt):
    caption = get_related_caption(prompt)
    generated_prompt = generate_prompt(caption)
    response = openai.Image.create(
      prompt=generated_prompt,
      n=1,
      size="256x256"
    )
    return caption, generated_prompt, response['data'][0]['url']

with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            with gr.Column():
                prompt_text = gr.Textbox(lines=5, label="Prompt")
                prompt_examples = gr.Examples(
                            examples=[
                                "きのこの山",
                                "たけのこの里",
                                "お菓子の家",
                            ],
                            inputs=[prompt_text],
                            outputs=None,
                        )
                btn = gr.Button(value="Generate Image")

            with gr.Column():
                caption = gr.Textbox(lines=5, label="Related Caption")
                generated_prompt = gr.Textbox(lines=5, label="Generated Prompt")
                out_image = gr.components.Image(type="filepath", label="Generated Image")

    btn.click(generate, inputs=[prompt_text], outputs=[caption, generated_prompt, out_image])
    demo.load()

demo.launch(share=True, auth=(app_username, app_password))