import os
import gradio as gr
import google.generativeai as genai

# APIキーの設定
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
model = genai.GenerativeModel(model_name='gemini-2.0-flash')

# 応答生成関数
def generate_response(message, history, temperature, top_p, top_k, max_output_tokens):
    # Gemini用の履歴に変換
    gemini_history = []
    for user, bot in history:
        gemini_history.append({"role": "user", "parts": [user]})
        gemini_history.append({"role": "model", "parts": [bot]})
    gemini_history.append({"role": "user", "parts": [message]})

    # モデルから応答を取得
    response = model.generate_content(
        gemini_history,
        generation_config={
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "max_output_tokens": int(max_output_tokens),
        }
    )

    # 履歴を更新して返す
    history.append((message, response.text))
    return "", history, history

# Gradio UI の構築
with gr.Blocks(theme='Nymbo/Alyx_theme') as demo:
    gr.Markdown("## Gemini Chatbot - Gemini 2.0 Flash + 調整可能パラメータ")

    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="メッセージを入力...")
    state = gr.State([])  # チャット履歴保持用

    # パラメータ調整スライダー
    with gr.Row():
        temperature = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Temperature")
        top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p")
        top_k = gr.Slider(1, 100, value=40, step=1, label="Top-k")
        max_output_tokens = gr.Number(value=1024, label="Max Output Tokens", precision=0)

    # 会話のリセットボタン
    clear = gr.Button("会話をリセット")

    # 入力時の処理
    msg.submit(
        generate_response,
        inputs=[msg, state, temperature, top_p, top_k, max_output_tokens],
        outputs=[msg, chatbot, state]
    )

    # クリア処理
    clear.click(lambda: ([], []), None, outputs=[chatbot, state])

# アプリケーションの起動
demo.launch()