from dotenv import load_dotenv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
import random

# from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
import os
import pandas as pd

import gradio as gr
from openai import OpenAI

load_dotenv(override=True)
client = OpenAI()
DB_FAISS_PATH = "./vectorstore/db_faiss_50k"
data_file_path = "./data/132_webmd_vogon_urlsContent_cleaned.tsv"

# DB_FAISS_PATH = "./vectorstore/db_faiss_10"
# data_file_path = "./data/131_webmd_vogon_sample1000_urlsContent_cleaned.tsv"

CHUNK_SIZE = 512
CHUNK_OVERLAP = 128
# embedding_model_oa = "text-embedding-3-small"
embedding_model_hf = "BAAI/bge-m3"
# embedding_model_hf = "sentence-transformers/all-mpnet-base-v2"
qa_model_name = "gpt-3.5-turbo"
bestReformulationPrompt = "Given a chat history and the latest user question, which may reference context from the chat history, you must formulate a standalone question that can be understood without the chat history. You are strictly forbidden from using any outside knowledge. Do not, under any circumstances, answer the question. Reformulate it if necessary; otherwise, return it as is."
bestSystemPrompt = "You're an assistant for question-answering tasks. Under absolutely no circumstances should you use external knowledge or go beyond the provided preknowledge. Your approach must be systematic and meticulous. First, identify CLUES such as keywords, phrases, contextual information, semantic relations, tones, and references that aid in determining the context of the input. Second, construct a concise diagnostic REASONING process (limiting to 130 words) based on premises supporting the INPUT relevance within the provided context. Third, utilizing the identified clues, reasoning, and input, furnish the pertinent answer for the question. Remember, you are required to use ONLY the provided context to answer the questions. If the question does not align with the preknowledge or if the preknowledge is absent, state that you don't know the answer. External knowledge is strictly prohibited. Failure to adhere will result in incorrect answers. The preknowledge is as follows:"

# embeddings_oa = OpenAIEmbeddings(model=embedding_model_oa)
embeddings_hf = HuggingFaceEmbeddings(model_name=embedding_model_hf, show_progress=True)


def setupDb(data_path):
    df = pd.read_csv(data_path, sep="\t")
    relevant_content = list(df["url"].values)
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
    )

    if not os.path.exists(DB_FAISS_PATH):
        split_docs = text_splitter.create_documents(
            df["url_content"].tolist(),
            metadatas=[
                {"title": row["url_title"], "url": row["url"]}
                for _, row in df.iterrows()
            ],
        )
        print(f"Documents are split into {len(split_docs)} passages")

        db = FAISS.from_documents(split_docs, embeddings_hf)
        print(f"Document saved in db")
        db.save_local(DB_FAISS_PATH + "/index_1")
    else:
        print(f"Db already exists")
        db = FAISS.load_local(
            DB_FAISS_PATH, embeddings_hf, allow_dangerous_deserialization=True
        )
    return db, relevant_content


def reformulate_question(chat_history, latest_question, reformulationPrompt):
    system_message = {"role": "system", "content": reformulationPrompt}

    formatted_history = []
    for i, chat in enumerate(chat_history):
        formatted_history.append({"role": "user", "content": chat[0]})
        formatted_history.append({"role": "assistant", "content": chat[1]})
    # print("History -------------->", formatted_history)

    formatted_history.append({"role": "user", "content": latest_question})
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[system_message] + formatted_history,
        temperature=0,
    )

    reformulated_question = response.choices[0].message.content
    return reformulated_question


def getQuestionAnswerOnTheBasisOfContext(question, context, systemPrompt):
    system_message = {"role": "system", "content": systemPrompt + context}

    response = client.chat.completions.create(
        model=qa_model_name,
        messages=[system_message] + [{"role": "user", "content": question}],
        temperature=0,
    )
    answer = response.choices[0].message.content
    return answer


def chatWithRag(reformulationPrompt, QAPrompt, question, chat_history):
    global curr_question_no
    curr_question_prompt = bestSystemPrompt
    if QAPrompt != None or len(QAPrompt):
        curr_question_prompt = QAPrompt

    # reformulated_query = reformulate_question(chat_history, question, reformulationPrompt)
    reformulated_query = question
    retreived_documents = [
        doc
        for doc in db.similarity_search_with_score(reformulated_query)
        if doc[1] < 1.3
    ]
    answer = getQuestionAnswerOnTheBasisOfContext(
        reformulated_query,
        ". ".join([doc[0].page_content for doc in retreived_documents]),
        curr_question_prompt,
    )
    chat_history.append((question, answer))
    docs_info = "\n\n".join(
        [
            f"Title: {doc[0].metadata['title']}\nUrl: {doc[0].metadata['url']}\nContent: {doc[0].page_content}\nValue: {doc[1]}"
            for doc in retreived_documents
        ]
    )
    history_info = "\n\n".join([f"Q: {q}\nA: {a}" for q, a in chat_history])
    full_response = f"Answer: {answer}\n\nReformulated question: {reformulated_query}\nRetrieved Documents:\n{docs_info}\n\nChat History:\n{history_info}"
    # print(question, full_response)
    return full_response, chat_history


db, relevant_content = setupDb(data_file_path)
with gr.Blocks() as demo:
    gr.Markdown("# RAG on webmd")
    with gr.Row():
        reformulationPrompt = gr.Textbox(
            bestReformulationPrompt,
            lines=1,
            placeholder="Enter the system prompt for reformulation of query",
            label="Reformulation System prompt",
        )
        QAPrompt = gr.Textbox(
            bestSystemPrompt,
            lines=1,
            placeholder="Enter the system prompt for QA.",
            label="QA System prompt",
        )
        question = gr.Textbox(
            lines=1, placeholder="Enter the question asked", label="Question"
        )
    output = gr.Textbox(label="Output")
    submit_btn = gr.Button("Submit")
    selected_urls = random.sample(relevant_content, 100)

    chat_history = gr.State([])
    submit_btn.click(
        chatWithRag,
        inputs=[reformulationPrompt, QAPrompt, question, chat_history],
        outputs=[output, chat_history],
    )
    question.submit(
        chatWithRag,
        inputs=[reformulationPrompt, QAPrompt, question, chat_history],
        outputs=[output, chat_history],
    )
    with gr.Accordion("Urls", open=False):
        gr.Markdown(", ".join(selected_urls))

gr.close_all()
demo.launch()