__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

# DATABASES = {
#     'default': {
#         'ENGINE': 'django.db.backends.sqlite3',
#         'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
#     }
# }
import streamlit as st
from huggingface_hub import InferenceClient
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, ServiceContext, PromptTemplate
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
import chromadb
from langchain.memory import ConversationBufferMemory
import pandas as pd
from langchain.schema import Document


# Set page config
st.set_page_config(page_title="MBAL Chatbot", page_icon="🛡️", layout="wide")

# Set your Hugging Face token here

HF_TOKEN = st.secrets["HF_TOKEN"]

@st.cache_resource
def init_chroma():
    persist_directory = "chroma_db"
    chroma_client = chromadb.PersistentClient(path=persist_directory)
    chroma_collection = chroma_client.get_or_create_collection("my_collection")
    return chroma_client, chroma_collection

@st.cache_resource
def init_vectorstore():
    persist_directory = "chroma_db"
    embeddings = HuggingFaceEmbeddings()
    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings, collection_name="my_collection")
    return vectorstore
@st.cache_resource
def setup_vector():
    # Đọc dữ liệu từ file Excel
    df = pd.read_excel("chunk_metadata_template.xlsx")
    chunks = []

    # Tạo danh sách các Document có metadata
    for _, row in df.iterrows():
        chunk_with_metadata = Document(
            page_content=row['page_content'],
            metadata={
                'chunk_id': row['chunk_id'],
                'document_title': row['document_title'],
                'topic': row['topic'],
                'access': row['access']
            }
        )
        chunks.append(chunk_with_metadata)

    # Khởi tạo embedding
    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')

    # Khởi tạo hoặc ghi vào vectorstore đã tồn tại
    persist_directory = "chroma_db"
    collection_name = "my_collection"

    # Tạo vectorstore từ dữ liệu và ghi vào Chroma
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=persist_directory,
        collection_name=collection_name
    )
    
    # Ghi xuống đĩa để đảm bảo dữ liệu được lưu
    vectorstore.persist()

    return vectorstore

# Initialize components
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=HF_TOKEN)
chroma_client, chroma_collection = init_chroma()
init_vectorstore()
vectorstore = setup_vector()

# Initialize memory buffer
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

def rag_query(query):
    # Lấy tài liệu liên quan
    retrieved_docs = vectorstore.similarity_search(query, k=5)
    context = "\n".join([doc.page_content for doc in retrieved_docs]) if retrieved_docs else ""

    # Lấy tương tác cũ
    past_interactions = memory.load_memory_variables({})[memory.memory_key]
    context_with_memory = f"{context}\n\nConversation History:\n{past_interactions}"

    # Chuẩn bị prompt
    messages = [
        {
            "role": "user",
            "content": f"""Bạn là một chuyên gia tư vấn, hỗ trợ khách hàng lựa chọn các sản phẩm bảo hiểm của MB Ageas Life tại Việt Nam. Vui lòng phản hồi một cách chuyên nghiệp và chính xác, đồng thời đề xuất các sản phẩm phù hợp bằng cách đặt một vài câu hỏi để tìm hiểu nhu cầu của khách hàng. Mọi thông tin cung cấp phải nằm trong phạm vi các sản phẩm của MB Ageas Life. Hãy mời khách hàng đăng ký để được tư vấn chi tiết hơn tại: https://www.mbageas.life/
{context_with_memory}
Câu hỏi: {query}
Câu trả lời:"""
        }
    ]

    response_content = client.chat_completion(messages=messages, max_tokens=1024, stream=False)
    response = response_content.choices[0].message.content.split("Answer:")[-1].strip()
    return response


def process_feedback(query, response, feedback):
   # st.write(f"Feedback received: {'👍' if feedback else '👎'} for query: {query}")
    if feedback:
        # If thumbs up, store the response in memory buffer
        memory.chat_memory.add_ai_message(response)
    else:
        # If thumbs down, remove the response from memory buffer and regenerate the response
       # memory.chat_memory.messages = [msg for msg in memory.chat_memory.messages if msg.get("content") != response]
        new_query=f"{query}. Tạo câu trả lời đúng với câu hỏi"
        new_response = rag_query(new_query)
        st.markdown(new_response)
        memory.chat_memory.add_ai_message(new_response)

# Streamlit interface

st.title("Chào mừng bạn đã đến với MBAL Chatbot")
st.markdown("***")
st.info('''
        Tôi sẽ giải đáp các thắc mắc của bạn liên quan đến các sản phẩm bảo hiểm nhân thọ của MB Ageas Life''')

col1, col2 = st.columns(2)

with col1:
    chat = st.button("Chat")
    if chat:
        st.switch_page("pages/chatbot.py")

with col2:
    rag = st.button("Store Document")
    if rag:
        st.switch_page("pages/management.py")

st.markdown("<div style='text-align:center;'></div>", unsafe_allow_html=True)