import gradio as gr
import time
from gemini_api import generate_query_with_gemini
from faiss_search import initialize_faiss, search_similar_documents
from data_loader import load_data_from_folder

folder_path = 'Data'
api_key = 'AIzaSyD3Hu747dbztC-jogggDfZudh_zYg40PJg'

model, index, documents = initialize_faiss(folder_path)

# Khởi tạo bộ nhớ ngữ cảnh để lưu giữ hội thoại trước
context_memory = []

# def generate_response_stream(chat_history, user_input):
#     global context_memory
#     try:
#         # Gọi API Gemini với ngữ cảnh hiện tại
#         sql_query = generate_query_with_gemini(user_input, api_key, model, index, documents, context_memory)
        
#         # Cập nhật lịch sử hội thoại
#         chat_history.append((user_input, sql_query))
#         yield gr.update(value=chat_history)
        
#         # Thêm câu hỏi và câu trả lời vào bộ nhớ ngữ cảnh
#         context_memory.append((user_input, sql_query))
#         # Giữ lại chỉ 3-5 vòng hội thoại gần nhất
#         if len(context_memory) > 5:
#             context_memory.pop(0)
            
#     except Exception as e:
#         chat_history.append(("Error", f"Error: {str(e)}"))
#         yield gr.update(value=chat_history)

def generate_response_stream(chat_history, user_input):
    global context_memory
    try:
        # Thêm câu hỏi của người dùng vào context_memory
        context_memory.append({"role": "user", "content": user_input})
        
        # Gộp toàn bộ context_memory thành một chuỗi văn bản để truyền vào truy vấn
        combined_context = "\n".join([f"{msg['role']}: {msg['content']}" for msg in context_memory])

        # Gửi truy vấn bao gồm bối cảnh đến Gemini API
        sql_query = generate_query_with_gemini(combined_context, api_key, model, index, documents)
        
        # Thêm phản hồi của chatbot vào cả chat_history và context_memory
        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": sql_query})
        
        # Lưu phản hồi của chatbot vào context_memory để làm bối cảnh cho tương tác sau
        context_memory.append({"role": "assistant", "content": sql_query})
        
        yield gr.update(value=chat_history)
    except Exception as e:
        chat_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
        yield gr.update(value=chat_history)

def clear_textbox():
    return ""

# with gr.Blocks() as iface:
#     gr.Markdown("<h1 id='title'>Hệ Thống Hỏi Đáp Y Tế VSS AI</h1>")
#     chatbot = gr.Chatbot(label="Trò chuyện", show_copy_button=True)
#     user_input = gr.Textbox(label="Nhập câu hỏi của bạn")
#     submit_button = gr.Button("Gửi câu hỏi")

#     user_input.submit(generate_response_stream, [chatbot, user_input], [chatbot]).then(clear_textbox, None, user_input)
#     submit_button.click(generate_response_stream, [chatbot, user_input], [chatbot]).then(clear_textbox, None, user_input)

# iface.launch(share=True)

with gr.Blocks() as iface:
    gr.Markdown("<h1 id='title'>Hệ Thống Hỏi Đáp Y Tế VSS AI</h1>")
    # Thêm `type="messages"` để tránh cảnh báo
    chatbot = gr.Chatbot(label="Trò chuyện", show_copy_button=True, type="messages")
    user_input = gr.Textbox(label="Nhập câu hỏi của bạn")
    submit_button = gr.Button("Gửi câu hỏi")

    user_input.submit(generate_response_stream, [chatbot, user_input], [chatbot]).then(clear_textbox, None, user_input)
    submit_button.click(generate_response_stream, [chatbot, user_input], [chatbot]).then(clear_textbox, None, user_input)

iface.launch(share=True)