import streamlit as st
import google.generativeai as genai
import os
import torch
import sys
import pysqlite3
from FlagEmbedding import FlagModel
# 將 sqlite3 指向 pysqlite3
sys.modules['sqlite3'] = pysqlite3

import chromadb
if hasattr(torch.classes, '__path__'):
    torch.classes.__path__ = []

########################################
# 設定 Google Gemini API & 主要參數
########################################
GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
genai.configure(api_key=GEMINI_API_KEY)

########################################
# 建立模型 (Gemini 2.0 Flash)
########################################
geminiModel = genai.GenerativeModel("gemini-2.0-flash")
checkModel = genai.GenerativeModel("gemini-2.0-flash")

########################################
# 判斷是否需要 RAG 查詢
########################################
def need_rag_check(model, user_input):
    prompt = f"""
你是一個智慧客服助理，專精於保險專業資訊。
無論使用者是詢問保險方案推薦、投保流程或其他與保險相關的專業問題，只要提到保險相關資訊，就必須查詢保險資料庫來輔助回答。
請根據下面的對話內容，僅輸出 "YES" ，表示需要查詢資料庫；否則輸出 "NO"。
對話內容：{user_input}
"""
    resp = model.generate_content(prompt)
    return ("YES" in resp.text.upper())


########################################
# 載入 Chuxin Embedding 模型 (快取)
########################################
@st.cache_resource
def load_embedding_model():
    return FlagModel(
        'chuxin-llm/Chuxin-Embedding',
        query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
        use_fp16=True
    )

embedding_model = load_embedding_model()

########################################
# 取得向量
########################################
def get_chuxin_embedding(text: str):
    return embedding_model.encode([text])[0].tolist()

########################################
# 初始化 ChromaDB
########################################
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(name="insurance_database", embedding_function=None)

########################################
# Streamlit 介面配置
########################################
st.title("💬 InsureGPT保險顧問")
st.write("🔍 與 AI 互動，詢問保險相關問題。")

# 自訂背景
st.markdown(
    """
    <style>
        .stApp {
            background: linear-gradient(
                rgba(238, 100, 2, 0.5),
                rgba(238, 100, 2, 0.5)
            ), url("https://imgur.com/LvZQr6p.png") no-repeat center center fixed;
            background-size: 50%;
            background-color: #fdd692;
            position: fixed;
            top: 0;
            left: 0;
            width: 100%;
            height: 100%;
            z-index: -1;
        }
    </style>
    """,
    unsafe_allow_html=True
)

########################################
# 建立 Chat Session 狀態
########################################
if "chat" not in st.session_state:
    # 初始化對話 (System Prompt)
    st.session_state.chat = geminiModel.start_chat(history=[
        {
            "role": "model",
            "parts": [(
                "背景設定: 你是一位精通台灣保險法規與保險商品的專業顧問，名字叫做InsureGPT。"
                "你可以使用知識庫中儲存的保險內容作為專業知識，回答使用者問題。"
                "也能根據使用者提供的個人資訊推薦客製化保險方案。"
                "若資訊不足或無法找到答案，請直接說明無法回答。"
                "請使用條列式或表格回答問題，以確保清楚易讀。"
                "請勿捏造資訊，如無法回答，請提供合理建議。"
                "在必要時，示範如何計算並提供適合的保險方案與試算金額。"
            )]
        }
    ])

########################################
# 用戶輸入 UI
########################################
user_input = st.text_input("💡 你的問題（例如：什麼是強制險？）", key="user_input")

########################################
# 按鈕事件
########################################
if st.button("送出"):
    if user_input.strip():
        # 先判斷是否需要 RAG
        do_rag = need_rag_check(checkModel, user_input)

        if do_rag:
            # Embedding + Query
            emb = get_chuxin_embedding(user_input)
            results = collection.query(query_embeddings=[emb], n_results=10)
            if results["documents"]:
                retrieved_docs = "\n\n---\n\n".join(results["documents"][0])
            else:
                retrieved_docs = "無匹配資料"
        else:
            retrieved_docs = "（模型判定此次提問無需檢索）"
        print(retrieved_docs)
        # 組合對話
        user_message = (
            f"你擁有以下與『{user_input}』相關的專業知識：\n\n"
            f"{retrieved_docs}\n\n"
            f"請回答下列問題：\n"
            f"{user_input}"
        )

        # 發送給 Gemini Chat
        chat = st.session_state.chat
        response = chat.send_message(user_message)

        # 顯示回應
        st.markdown(f"### 🤖 InsureGPT 回應：\n{response.text}")

        # 更新 chat 狀態
        st.session_state.chat = chat
        st.session_state.input_value = "" 
    else:
        st.warning("請輸入問題！")

########################################
# 顯示聊天紀錄
########################################
st.subheader("💬 聊天記錄")
for msg in st.session_state.chat.history:
    role = "user" if msg.role == "user" else "assistant"
    if role == "assistant" and "背景設定" in msg.parts[0].text:
        continue
    if role == "user":
        question_start = msg.parts[0].text.find("問題：\n")
        if question_start != -1:
            msg_text = msg.parts[0].text[question_start:]
        else:
            msg_text = msg.parts[0].text
    else:
        msg_text = msg.parts[0].text
    with st.chat_message(role):
        st.markdown(msg_text)