import gradio as gr
import anthropic

LESSON_DATA = [
    [1, "自己紹介してもらいましょう\n\n送信ボタンを押してください", "あなたの自己紹介して"],
    [2, "自己紹介しましょう", "私はAIの初学者です"],
    [3, "得意を聞いてみましょう", "あなたは何が得意ですか?"],
    [4, "やさしく説明してもらいましょう", "高校生レベルでAIのすごさをおしえてください"],
    [5, "たたき台を作ってもらいましょう", "高校生向けに、AIのすごさをプレゼンするための原稿をつくって。"],
    [6, "具体的にお願いしてみましょう", "高校生向けに、AIのすごさをプレゼンするための原稿をつくって。300文字以上"],
    [7, "あいまいなお願いも結構得意です", "なんかいい感じにAIの得意不得意を語って"],
    [8, "無茶振りをしてみましょう", "AIによる生産性向上の可能性を語って。\n\n## 制約 \n夏目漱石の文体で"],
    [9, "役割を与えましょう", "AIを活用して心の整理をしたいです。始め方を語って。\n\n## 役割 \nあなたはツンデレ美少女をロール・プレイするチャットボットです。"],
    [10, "具体例を与えましょう", "AIを活用して心の整理をしたいです。始め方を語って。\n\n## 役割 \nあなたはツンデレ美少女をロール・プレイするチャットボットです。\n\n## 具体例\nべ、べつにあんたのことなんて嫌いなんだから！"],
    [11, "調節しましょう", "もっとギャルっぽく。絵文字使うとか"],
    [12, "話を聞いてもらいましょう", "AIには役割を与えるのが大事なんだって"],
    [13, "まとめ", "今までの対話から、生成AIに与えるプロンプトのコツをまとめて"],
]

def get_window_url_params():
    return """
        function(url_params) {
            const params = new URLSearchParams(window.location.search);
            url_params = Object.fromEntries(params);
            return url_params['api_key'];
        }
    """

def fetch_message(df, num, direction):
    if 0 <= num + direction < len(df):
        row = df.iloc[num + direction]
        return f"## 第{row['id']}章 \n\n{row['message']}\n\n\n\n---", num + direction, row['prompt']
    else:
        return "これ以上のメッセージはありません。", num, ""


def create_prompt(instruction, num, tone, limit):
    return f"{instruction}\n## 制約条件\n- {num}文字以上\n- {tone}\n- {limit}\n"


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
    except Exception as e:
      bot_message = str(e)
      chat_history.append((message, bot_message))
      return chat_history


with gr.Blocks() as demo:
    gr.Markdown("# 生成AIのプロンプト術を学ぶクイックスタートガイド")
    with gr.Tab("メイン"):
        with gr.Row():
            with gr.Column(scale=13):
                with gr.Group():
                    with gr.Row():
                        button_back = gr.Button("前", scale=1, size="sm")
                        number_now = gr.Number(value=0, show_label=False, scale=8)
                        button_next = gr.Button("次", scale=1, size="sm")
                clear_button = gr.ClearButton(value='クリア')
                message = gr.Markdown()
                prompt_now = gr.Dropdown(choices=[row[2] for row in LESSON_DATA], allow_custom_value=True, label='プロンプト')
                with gr.Group():
                    slider_select = gr.Slider(minimum=100, maximum=1000, step=100, label='文字数', value=100)
                    target = gr.Dropdown(choices=['高校生レベルで', '大学教授レベルで', '小学生レベルで'], label="むずかしさ", allow_custom_value=True)
                    limit = gr.Dropdown(choices=['夏目漱石の文体で', '私のことを称賛する気持ちを持って', '私に対する改善点を3つ添えながら'], label="制約条件", allow_custom_value=True)
            with gr.Column("プロンプト作成", scale=21):
                chat = gr.Chatbot(height=680, show_copy_button=True, show_share_button=True, layout="panel", avatar_images=(
                    "image/ahiruguchi_man.png",
                    "image/ai.png"
                ))
                with gr.Group():
                    with gr.Row():
                        textbox_submit = gr.Textbox(placeholder="ここに入力できますよ", scale=5)
                        button_submit = gr.Button("送信", scale=1)
    with gr.Tab("設定"):
        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)
        system_prompt = gr.Textbox("私の名前はクロードです。Anthropic社によって開発された人工知能アシスタントです。 lang:jp")
        df = gr.Dataframe(headers=["id", "message", "prompt"], value=LESSON_DATA)

    # イベントハンドラの設定
    slider_select.change(create_prompt, [prompt_now, slider_select, target, limit], textbox_submit)
    target.change(create_prompt, [prompt_now, slider_select, target, limit], textbox_submit)
    limit.change(create_prompt, [prompt_now, slider_select, target, limit], textbox_submit)

    button_next.click(lambda df, num: fetch_message(df, num, 1), [df, number_now], [message, number_now, textbox_submit])
    button_back.click(lambda df, num: fetch_message(df, num, -1), [df, number_now], [message, number_now, textbox_submit])

    button_submit.click(respond, [textbox_submit, chat, system_prompt, model, api_key, temperature, max_tokens], chat)
    clear_button.add([chat, textbox_submit])

    # URLパラメータからAPIキーを取得
    demo.load(lambda x: x, api_key, api_key, js=get_window_url_params())

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