import streamlit as st
import os
from groq import Groq
import re

# コードブロックを抽出し、プレースホルダーに置き換える関数
def extract_code_blocks(text):
    code_blocks = re.findall(r'```.*?```', text, re.DOTALL)
    for i, block in enumerate(code_blocks):
        placeholder = f"<CODE_BLOCK_{i+1}>"
        text = text.replace(block, placeholder)
    return text, code_blocks

# プレースホルダーを実際のコードブロックに置き換える関数
def reinsert_code_blocks(translated_text, code_blocks):
    for i, block in enumerate(code_blocks):
        placeholder = f"<CODE_BLOCK_{i+1}>"
        translated_text = translated_text.replace(placeholder, block)
    return translated_text

# APIキーの設定
try:
    client = Groq(api_key=st.secrets["GROQ_API_KEY"])
    st.success("APIクライアントが正常に設定されました。")
except Exception as e:
    st.error(f"APIクライアントの設定に失敗しました: {e}")

# Streamlitアプリケーションのタイトルと説明
st.title("Groq APIによるAIチャットボット")
st.subheader("生成されたテキストがリアルタイムで表示されます")

# 会話履歴の初期化（セッションごとに維持）
if 'conversation_history' not in st.session_state:
    st.session_state['conversation_history'] = []

# 過去の会話履歴を表示
for entry in st.session_state['conversation_history']:
    st.chat_message(entry['role']).markdown(entry['content'])

# ユーザー入力の受付
if prompt := st.chat_input("質問を入力してください:"):
    st.session_state['conversation_history'].append({"role": "user", "content": prompt})
    st.chat_message("user").markdown(prompt)
    
    translated_prompt = ""
    response_text = ""
    translated_response_text = ""
    
    with st.chat_message("assistant"):
        message_placeholder = st.empty()  # 応答の一時表示用の空要素
        
        try:
            # 日本語の質問を英語に翻訳（gemma2-9b-itを使用）
            translation_response = client.chat.completions.create(
                model="gemma2-9b-it",
                messages=[
                    {"role": "system", "content": "Translate the following Japanese text to English accurately without adding or omitting any details."},
                    {"role": "user", "content": prompt}
                ],
                stream=True,
                temperature=0.8,
                top_p=0.92
            )
            
            # 英語に翻訳された質問を蓄積
            for chunk in translation_response:
                if hasattr(chunk.choices[0].delta, 'content'):
                    chunk_text = chunk.choices[0].delta.content
                    if chunk_text:
                        translated_prompt += chunk_text

            # 英語に翻訳された質問をllama3-70b-8192で応答生成
            response = client.chat.completions.create(
                model="llama3-70b-8192",
                messages=[
                    {"role": "system", "content": "Provide a detailed response including code examples."},
                    {"role": "user", "content": translated_prompt}
                ],
                stream=True
            )
            
            # 応答の処理とコードブロックの抽出
            full_response_text = ""
            for chunk in response:
                if hasattr(chunk.choices[0].delta, 'content'):
                    chunk_text = chunk.choices[0].delta.content
                    if chunk_text:
                        full_response_text += chunk_text
            
            # コードブロックを抽出し、プレースホルダーを挿入
            text_with_placeholders, code_blocks = extract_code_blocks(full_response_text)

            # 非コード部分を日本語に翻訳
            final_translation_response = client.chat.completions.create(
                model="gemma2-9b-it",
                messages=[
                    {"role": "system", "content": "Translate the following English text to Japanese, keeping all code blocks in their original form."},
                    {"role": "user", "content": text_with_placeholders}
                ],
                stream=True,
                temperature=0.8,
                top_p=0.92
            )
            
            for chunk in final_translation_response:
                if hasattr(chunk.choices[0].delta, 'content'):
                    chunk_text = chunk.choices[0].delta.content
                    if chunk_text:
                        translated_response_text += chunk_text
            
            # コードブロックをプレースホルダーの位置に再挿入
            final_output = reinsert_code_blocks(translated_response_text, code_blocks)
            message_placeholder.markdown(final_output)
            st.session_state['conversation_history'].append({"role": "assistant", "content": final_output})
            st.success("APIからの応答が正常に処理されました。")
        
        except Exception as e:
            error_message = f"APIリクエストの処理中にエラーが発生しました: {e}"
            st.error(error_message)
            st.write("Error:", error_message)  # エラーの表示

# 会話履歴をリセットするボタン
if st.button("会話をリセット"):
    st.session_state['conversation_history'] = []
    st.rerun()

# デバッグ情報の表示（開発中のみ使用）
if st.checkbox("デバッグ情報を表示"):
    st.write("会話履歴:", st.session_state['conversation_history'])
    st.write("API Key設定状況:", "設定済み" if "GROQ_API_KEY" in st.secrets else "未設定")