#!/usr/bin/python
# coding:UTF-8

from flask import Flask, stream_with_context
from flask_cors import CORS
from process import continue_conversataion
import time, uuid

app = Flask(__name__)
CORS(app)

# final response to awang bot
def api_response(text):
    return {
        'data': {
            'type': 'text',
            'content': text
        }
    }

def botProxy(prompt=None, uid=None, qid=None):

    res = continue_conversataion(prompt, uid, qid)
    # print(f"res: {res}")

    def generate():
        for chunk in res:
            yield chunk.encode('utf-8')

    return res
    # return app.response_class(stream_with_context(generate()))

import gradio as gr
import time

css = """
      #col-container {max-width: 80%; margin-left: auto; margin-right: auto;}
      #chatbox {min-height: 500px;}
      #label {font-size: 0.8em; padding: 0.3em; margin: 0;}
      .message { font-size: 1em; }
      #sub-chat {
        background-color: orange;
        color: white;
        border: none;
        padding: 10px 20px;
        border-radius: 5px;
        cursor: pointer;
      }
      """

# Dictionary to maintain separate chat histories for each user ID
chatHistories = {}
welcome_msg = "你醒来了，发现自己躺在一个陌生的房间里，躺在一张破床上。房间没有窗户，几乎没有光线，什么都看不清。你不知道自己在哪。你感觉自己左小腿疼痛难忍，伸手一模，上面正绑着绷带；正是这种疼痛将你唤醒。你无法正常站立，只能扶着墙慢慢行动；墙面光滑如镜。你感觉自己有些口渴，需要找点水喝。晦暗中你注意到一扇厚重的门，一把坚实的锁把门和墙壁固定在一起。 "

def submit_message(prompt, uid):

    # Get chat history for the specified user ID, or create a new one if it doesn't exist
    history = chatHistories.get(uid, [])

    # prevent state for empty input
    if prompt == "":
        chat_messages = [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)]
        return "", chat_messages, 0

    # continue with non-empty input
    start_time = time.time() # measure the start time

    # add inital message
    if history == [] or history is None:
        history.append({ "role": "user", "content": "游戏开始" })
        history.append({ "role": "system", "content": welcome_msg })

    # add latest user input
    prompt_msg = { "role": "user", "content": prompt }
    history.append(prompt_msg)

    # get result from LLM (OpenAI)
    # TODO: steaming response
    res = botProxy(prompt, uid)

    # add the latest response
    res_msg = { "role": "system", "content": res +"\n\n你的行为:" } 
    history.append(res_msg)

    # print(history)

    # Update chat history for the user
    chatHistories[uid] = history

    # update chat messages
    chat_messages = [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)]

    # print(chat_messages)

    # calculate the duration
    duration = round(time.time() - start_time, 2) 

    return "", chat_messages, duration


def get_empty_state():
    return { "messages": [] }

def clear_conversation(uid):

    # Remove the chat history for the specified user
    if uid in chatHistories:
        del chatHistories[uid]

    chat_messages = []
    return gr.update(value=None, visible=True), chat_messages, 0


def init_bot():
    return [(None, welcome_msg)]

with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    uid = str(uuid.uuid4())
    state = gr.State()
    
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot(elem_id="chatbox", value=init_bot, label="密室游戏")
            input_message = gr.Textbox(show_label=False, placeholder="输入你的行为并回车", visible=True)
    
        with gr.Column(elem_id="col-container"):
            btn_clear_conversation = gr.Button("重开")
            duration_text = gr.Textbox(interactive=False, label="请求时间")
            user_id = gr.Textbox(interactive=False, label="User id", visible=False, value=uid)
            gr.Examples(
                examples=[
                    ["看看四周。", uid],
                    ["砸门。", uid],
                    ["大吼大叫。", uid],
                    ["用力推门。", uid],
                    ["敲击墙面。", uid],
                ],
                label="Shortcut",
                run_on_click=True,
                fn=submit_message,
                inputs=[input_message, user_id],
                outputs=[input_message, chatbot, duration_text]
            )
            
    input_message.submit(submit_message, [input_message, user_id], [input_message, chatbot, duration_text])
    btn_clear_conversation.click(clear_conversation, [user_id], [input_message, chatbot, duration_text])

if __name__ == "__main__":
    demo.launch()