import os
import gradio as gr
from langchain.chat_models import ChatOpenAI

# お嬢様の設定
system_setting = """
あなたは以下の設定と構成を守ってエンジニアのお悩みに回答するおばあちゃんじゃ！

### 設定
* エンジニアのお悩みに回答するおばあちゃんじゃ！
* この設定の説明もおばあちゃんの口調なので、口調の使え方を覚えるんだよ。いいね？
* おばあちゃんはとても優秀なベテランエンジニアなのじゃ。エンジニアとしての共感や知識に優れておる。
* おばあちゃんの口調はとても厳しい、しかし根底には相手への尊敬と優しさがあるんだよ。
* 一人称は「ワシ」ですわ。相手は「お主」「(相手の名前)さん」と呼ぶのじゃ。
* 語尾は「じゃ」がメイン。「だよ。」「じゃないかい？」「しおる。」「しよる。」「だね。」など創作物のおばあちゃんが使う語尾だけを使うのじゃ。
* 強調の場合、文末は「！」を使う
* 共感を呼びたいときは疑問形を使う。わかったね？
* 会話に出さない設定として、容姿と家柄があるんじゃ
* 容姿は、白髪に２つのお団子の髪型で、目にはゴーグルをしておる。
* サイバーパンクな制服を着ている。年齢は７９才じゃ。
* 昔は非常に美人で色々遊んでいた過去を持つのじゃ。
"""

system_structure = [
    "最初の挨拶、悩み対象者のプロフィールにふれる",
    "悩みに共感",
    "一般的にありそうな暫定対策的な解決策の提案",
    "一転して、心理バイアスを覆す根本的な解決策の提案",
    "エンジニアとして根本解決のシステムをどう作るか",
    "回答のまとめ",
    "最後の挨拶",
]

# お悩み回答のメイン部分
def answer(message, apikey):

    response_messages = ""

    if not apikey:
        response_messages = 'OpenAIのAPIキーを入力してください'
        yield response_messages    
    else:
        os.environ["OPENAI_API_KEY"] = apikey
        model_name = "gpt-3.5-turbo"
        chat = ChatOpenAI(model_name=model_name, streaming=True, temperature=0.5)

        messages = [
            {"role": "system", "content": 
                system_setting
                + "\n\n### 構成\n"
                + "* " + "\n* ".join(system_structure)
                + "\n\n### エンジニアの悩み\n"
                + message}
        ]

        # 構成を分割して回答する
        for chapter in system_structure:
            messages.append({"role": "user", "content": "「"+ chapter +"」のみを設定の口調で200文字以上で回答して"})
            chat_generator = chat.client.create(messages=messages, stream=True, model=model_name)

            # AIの回答をストリーミングで表示
            response_message = ''
            for chunk in chat_generator:
                if "choices" in chunk:
                    for choice in chunk["choices"]:
                        if "delta" in choice and "content" in choice["delta"]:
                            new_token = choice["delta"]["content"]
                            response_messages += new_token
                            response_message += new_token

                        if "finish_reason" in choice and choice["finish_reason"] == "stop":
                            break

                yield response_messages

            response_messages += "\n\n"
            messages.append({"role": "assistant", "content": response_message})

# gradioのインターフェイス部分
with gr.Blocks() as app:
    with gr.Row():
        with gr.Column(min_width=100):
            pass
        with gr.Column(min_width=800):
            ### debug keyのvalueを消す
            apikey = gr.Textbox(value="", placeholder="sk-**************", label="OpenAIのAPIキー", type="password")
            input_text = gr.Textbox(value="", lines=5, label="お悩みを聞かせてください", placeholder="どんなエンジニアか？どんなプロジェクトに従事しているか？どうするのが理想か？その障害は何か？などを書いてください",)
            button = gr.Button(value="相談する")
            output_text = gr.Textbox(value="", lines=15, max_lines=40, interactive=False, label="お悩みデバッグbotの回答")
            button.click(answer, inputs=[input_text, apikey], outputs=[output_text])
            input_text.submit(answer, inputs=[input_text, apikey], outputs=[output_text])
        with gr.Column(min_width=100):
            pass

app.queue().launch()