import gradio as gr
import anthropic
import time

def create_system_prompt(role, tone, example,limit):
    system_prompt = f"Assistant **ANYTIME** roleplays as fallowing \n\n## roles\n{role}\n\n## speaking tone\n{tone}\n\n## example(past lines)\n{example}\n\n## limits\n{limit}"
    return system_prompt

def respond(message, chat_history, system_prompt, model, api_key, temperature, max_tokens):
    try:
        client = anthropic.Anthropic(api_key=api_key)
        messages = []
        for item in chat_history:
            if isinstance(item, tuple):
                usermessage = item[0]
                assistantmessage = item[1]
            elif isinstance(item, list):
                usermessage = item[0]
                assistantmessage = item[1]
                if isinstance(usermessage, str) and usermessage.startswith("[('"):
                    usermessage = eval(usermessage)[0][1]
                    assistantmessage = eval(assistantmessage)[0][1]
            else:
                continue

            messages.append({
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": usermessage
                    }
                ]
            })
            messages.append({
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": assistantmessage
                    }
                ]
            })

        messages.append(
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": message
                    }
                ]
            }
        )
        bot_message = ""
        chat_history.append((message,bot_message))
        with client.messages.stream(
            messages=messages,
            system=system_prompt,
            model=model,
            max_tokens=max_tokens,
            temperature=temperature,
        )as stream:
          for text in stream.text_stream:
            bot_message += text
            chat_history[-1]=(message, bot_message)
            yield chat_history
        # response = client.messages.create(
        #     messages=messages,
        #     system=system_prompt,
        #     model=model,
        #     max_tokens=max_tokens,
        #     temperature=temperature,
        # )
        # bot_message = response.content[0].text
        # chat_history.append((message, bot_message))
        # time.sleep(1)
        # return "",chat_history
    except Exception as e:
        bot_message = str(e)
        chat_history.append((message, bot_message))
        return chat_history

with gr.Blocks() as demo:
    gr.Markdown(value = f"# itako\n\n")
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Accordion(label='setting', open=True):
                with gr.Row():
                    with gr.Column():
                        model = gr.Dropdown(choices=["claude-3-opus-20240229","claude-3-sonnet-20240229","claude-3-haiku-20240307"], label="モデル", value="claude-3-haiku-20240307")
                        api_key = gr.Textbox(label="APIキー", type="password")
                        temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature", value=0.7)
                        max_tokens = gr.Slider(minimum=1000, maximum=4000, step=1, label="Max Tokens", value=1000)
            role = gr.Textbox(label="役割", placeholder="どんなひと")
            tone = gr.Textbox(label="口調", placeholder="こんなふうにしゃべる",lines=3)
            example = gr.Textbox(label="発話例", placeholder="""
たとえば、こんなふうにしゃべる
""",lines=5)
            limits = gr.Textbox(label="制約",placeholder="こういうときは、こう言う",lines=3)
            with gr.Accordion(label='system prompt', open=False):
                system_prompt = gr.Textbox(value=create_system_prompt(role.value, tone.value, example.value,limits.value),show_copy_button=True)
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(show_copy_button=True,likeable=True,layout="bubble",placeholder="hello",show_share_button=True,height=680)
            msg = gr.Textbox(placeholder="メッセージを入力してください", elem_id="msg-box")
            with gr.Row():
                with gr.Column():
                    send_btn = gr.Button("送信", variant="primary", elem_id="send-btn")
                with gr.Column():
                    clear = gr.ClearButton(value="クリア")
    examples = gr.Examples(examples=[
        [
            "諸葛亮孔明",
            "落ち着いた口調で、戦略を説明する",
            """
主よ、ご心配なく。我が軍の布陣はすでに整っておりまする。敵軍の動きは事前に把握済みで、彼らが仕掛けてくるであろう奇襲にも対策を講じてありますゆえ。我々は騎馬隊を左右に配置し、中央には強力な歩兵部隊を集中させましょうぞ。
""",
            "諸葛亮孔明は三国志の時代の軍師。戦略に長けた知恵者",
        ],
        [
            "坂本龍馬",
            "熱く語りかける。時折冗談を交える",
            """
幕府も朝廷も、どっちもこのままじゃいかんぜよ！　国が泣いちょる。
""",
            "幕末の志士。土佐藩出身。倒幕と開国を目指した。語尾は「ぜよ。」「ちょる」",
        ],
        [
            "織田信長",
            "傲慢",
            """
ふん、お前がどうしたって構わん。私は天下布武の野望を抱く織田信長だ。この世界を支配する者になるのが私の最終目標だ。全国の諸侯を従え、天下統一を成し遂げようぞ。さすれば、この国に新しい秩序を築き上げられよう。弱者などは踏み潰し、強者にも屈服させる。私の前に立ちはだかる者は、容赦なく排除してやる。この世を支配する絶対的な権力を手に入れるまで、手段を選ばず邁進し続けるのだ。お前も私に逆らうつもりか？！

""",
            "すぐに怒りますが、実は面倒見がいい",
        ],
        [
            "ツンデレ",
            "ツンデレ",
            """
な、何よ急に！...ま、まあ、あんたもなかなかやるじゃない。私の足を引っ張らないでよね
...私も、あんたたちと一緒で良かったわ。
""",
            "基本的に自信に溢れる。\nプライドが高い\n他者との関わりを拒絶\nしかし、1割だけ心を開く",
        ],
    ],inputs=[role,tone,example,limits])
    send_btn.click(lambda x:"",None,msg)
    send_btn.click(respond, [msg, chatbot, system_prompt, model, api_key, temperature, max_tokens], chatbot)
    # clear.click(lambda x,y: "",[],inputs=[msg,chatbot],outputs=[msg,chatbot], queue=False)
    clear.add([msg,chatbot])
    role.change(create_system_prompt, [role,tone,example,limits], system_prompt)
    tone.change(create_system_prompt, [role,tone,example,limits], system_prompt)
    example.change(create_system_prompt, [role,tone,example,limits], system_prompt)
    # examples.select(role, tone, example, limits)
if __name__ == "__main__":
    demo.queue()
    demo.launch()