import streamlit as st
import faiss
import numpy as np
import json
from sentence_transformers import SentenceTransformer
import time

# データを読み込む
with open("data/qa_data.json", "r", encoding="utf-8") as f:
    data = json.load(f)

questions = [item["question"] for item in data]
answers = [item["answer"] for item in data]

# 埋め込みモデルをロード
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja")

# FAISSインデックスをロード
index_q = faiss.read_index("data/faiss_question.index")
index_a = faiss.read_index("data/faiss_answer.index")

# サイドバー設定
st.set_page_config(initial_sidebar_state="collapsed")
with st.sidebar.expander("⚙️ 設定", expanded=False):
    threshold_q = st.slider("質問の類似度しきい値", 0.0, 1.0, 0.7, 0.01)
    threshold_a = st.slider("回答の類似度しきい値", 0.0, 1.0, 0.65, 0.01)


def search_answer(user_input):
    """FAISSを使用して最適な回答を検索"""
    user_embedding = model.encode([user_input]).astype(np.float32)

    # 質問に対して検索
    D_q, I_q = index_q.search(user_embedding, 1)
    score_q = 1 / (1 + D_q[0][0])

    if score_q >= threshold_q:
        # Replace \n with markdown line breaks
        return answers[I_q[0][0]].replace("\n", "  \n"), f"質問にマッチ ({score_q:.2f})"

    # 回答に対して検索
    D_a, I_a = index_a.search(user_embedding, 1)
    score_a = 1 / (1 + D_a[0][0])

    if score_a >= threshold_a:
        # Replace \n with markdown line breaks
        return answers[I_a[0][0]].replace("\n", "  \n"), f"回答にマッチ ({score_a:.2f})"

    return "申し訳ありませんが、ご質問の答えを見つけることができませんでした。もう少し詳しく説明していただけますか？", "一致なし"


def stream_response(response):
    """レスポンスをストリーム表示する（文字単位）"""
    for char in response:
        if char == "\n":
            # Replace newline with markdown line break
            yield "  \n"
        else:
            yield char
        time.sleep(0.05)


# Streamlitチャットインターフェース
st.title("🤖 日本語Q&Aチャットボット")

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if user_input := st.chat_input("質問を入力してください:"):
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    with st.spinner("考え中... お待ちください。"):
        answer, info = search_answer(user_input)

    with st.chat_message("assistant"):
        response_placeholder = st.empty()
        response_placeholder.write_stream(stream_response(answer))

    st.session_state.messages.append({"role": "assistant", "content": answer})