import os
import gradio as gr
import base64
from openai import OpenAI, BadRequestError

OUT_FOL = "output"

def create_image(image_path, request: gr.Request):

    err_msg = ""
    save_image_path = None

    prompt = """
    この画像をジブリ風にして下さい。
    """

    file_name = os.path.basename(image_path)

    save_file_name = "Ghi_" + file_name

    os.makedirs(OUT_FOL, exist_ok=True)

    save_image_path = OUT_FOL + "/" + save_file_name

    client = OpenAI()

    try:

        result = client.images.edit(
            model="gpt-image-1",
            image=open(image_path, "rb"),
            prompt=prompt,
            # quality="low",
            size="1024x1024"
        )

        image_base64 = result.data[0].b64_json
        image_bytes = base64.b64decode(image_base64)

        with open(save_image_path, "wb") as f:
            f.write(image_bytes)

    except BadRequestError as e:
        err_msg = "画像が受け付けられませんでした。別の画像で試してください。"
        print(e)
    except Exception as e:
        err_msg = "画像作成でエラーが発生しました。別の画像で試してください。"
        print(e)

    return save_image_path, err_msg

with gr.Blocks(theme=gr.themes.Ocean()) as demo:

    title = "<center><h2>ジブリ風イラストメーカー</h2></center>"
    message = "<center><h3><b>使い方は画像をアップして画像作成ボタンを押すだけです。<br>"
    message += "※1日ごとに利用制限があります。利用できない場合は翌日お試しください。"
    message += "</b></h3></center>"

    gr.Markdown(title)
    gr.Markdown(message)

    with gr.Row():

        with gr.Column():

          # 入力
          input_image = gr.Image(label="入力画像", type="filepath")

          btn = gr.Button("画像作成")

        with gr.Column():

          # 出力
          out_image = gr.Image(label="生成画像", type="filepath", interactive=False)

          sys_msg = gr.Text(label="システムメッセージ")

    btn.click(create_image, input_image, [out_image, sys_msg])

if __name__ == '__main__':
    demo.queue()
    demo.launch()