import gradio as gr
import os
from langchain.retrievers import EnsembleRetriever
from utils import *
import requests
from pyvi import ViTokenizer, ViPosTagger
import time
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain_community.chat_message_histories import ChatMessageHistory

retriever = load_the_embedding_retrieve(is_ready=False, k=3)
bm25_retriever = load_the_bm25_retrieve(k=3)

ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, retriever], weights=[0.2, 0.8]
)

llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0, openai_api_key= os.environ["OPENAI_API_KEY"])


def greet3(quote, history):
    print(history)
    demo_ephemeral_chat_history = ChatMessageHistory()
    if history == '':
        history = [("Bạn có thể giải thích về quy chế và quyền của sinh viên tại trường này không?",
                '''Quy chế và quyền của sinh viên tại trường Đại học Mở TP.HCM được quy định rõ trong các điều khoản sau:
                1. Hiệu trưởng Trường có quyền ra quyết định thành lập và quy định cụ thể về chức năng, nhiệm vụ, tổ chức và hoạt động của Hội đồng khen thưởng và kỷ luật sinh viên.
                2. Sinh viên có quyền khiếu nại về khen thưởng, kỷ luật. Khi có vi phạm kỷ luật, sinh viên có quyền được phân tích và đề nghị hình thức kỷ luật thông qua việc họp với các tổ chức sinh viên và gửi biên bản họp đến phòng Công tác sinh viên để trình Hội đồng.
                3. Sinh viên có quyền đề đạt nguyện vọng và khiếu nại lên Hiệu trưởng Trường để giải quyết các vấn đề có liên quan đến quyền, lợi ích chính đáng của sinh viên.
                4. Sinh viên được hỗ trợ giới thiệu nhà trọ theo quy định của trường.
                
                Các chủ đề liên quan mà bạn có thể muốn tìm hiểu thêm:
                - Quy chế và quyền của sinh viên tại các trường đại học khác.
                - Hệ thống hỗ trợ sinh viên tại trường Đại học Mở TP.HCM.
                - Quy trình khiếu nại và giải quyết tranh chấp sinh viên tại trường Đại học Mở TP.HCM.
                '''),
                ("Ai là Dương Hữu Thành?",
                "Dương Hữu Thành là Thạc sĩ và Phó Trưởng Bộ môn của Bộ môn Cơ sở ngành tại Khoa Công nghệ thông tin.")]
        for user, assistant in history:
            demo_ephemeral_chat_history.add_user_message(user)
            demo_ephemeral_chat_history.add_ai_message(assistant)
    else:
        for user, assistant in eval(history):
            demo_ephemeral_chat_history.add_user_message(user)
            demo_ephemeral_chat_history.add_ai_message(assistant)

    # Summary the message

    chat_history = summarize_messages(demo_ephemeral_chat_history=demo_ephemeral_chat_history, llm=llm).messages

    # Get the new question
    new_question = get_question_from_summarize(chat_history[0].content, quote, llm)

    # Retrieve
    documents_query = ensemble_retriever.invoke(new_question)

    print(documents_query)

    context = ''
    for i in documents_query:
        context += i.page_content + '\n'

    # Get answer

    answer = get_final_answer(question=new_question, context=context,
                              prompt=os.environ['PROMPT'], llm=llm)

    return new_question, answer


if __name__ == "__main__":
    iface = gr.Interface(fn=greet3, inputs=["text", "text"], outputs=["text", "text"])
    iface.launch()