import gradio as gr
import requests
import os
from PIL import Image
from io import BytesIO
import base64
import uuid

os.mkdir("/tmp/gradio")

def respond(message, chat_history):
  try:
    history_send=[(m,b) for m,b in chat_history if m is not None]
    response = requests.post(os.environ["BACKEND"], json={
        "data": [
            message,
            history_send
        ]
    }).json()

    #print(response)
      
    data = response["data"]
    #print(data)
    image=Image.open(BytesIO(base64.b64decode(data[0])), formats=["PNG"])

    image_path="/tmp/gradio/"+uuid.uuid4().hex+".png"
    image.save(image_path)

    chat_history.append((data[2][-1][0],data[2][-1][1]))
    chat_history.append((None, (image_path,)))
    chat_history.append((None, "上記プロンプトで画像を生成しました。画像を変更したい場合や新しく作りたい場合は教えてください。"))
      
    return "", chat_history
        
  except Exception as e:
    print(e)
    return "Error",chat_history
      
css = """
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
"""

title = """
<div style="text-align: center;max-width: 700px;">
    <h1>ChatEmi</h1>
    <p>個人情報や機密情報を入力しないでください。</p>
    <p>※プロンプトの出力が荒れるのは国産LLMの限界です。今後の活躍にご期待ください。</p>
</div>
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)
        chatbot = gr.Chatbot(label='対話履歴', avatar_images=("me.png", "emi.png"))
        msg = gr.Textbox(label="プロンプト欄", placeholder="生成したい画像を入力してエンターキーを押してください。")
        clear = gr.ClearButton([msg, chatbot])

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

    
demo.queue(concurrency_count=1)
demo.launch()