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_largeChunk"
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. You are strictly forbidden from adding extra things in the question if not required. Do not, under any circumstances, answer the question. Reformulate ONLY if it is necessary; otherwise, return it as is.

Example: I am getting fat, how can I loose weight?
Reformulated question: What are some effective strategies for losing weight in a healthy manner?
This reformulation is BAD since it added extra thing "in a health manner". The question was understandable even without chat history.

Example: When was he born
Chat History: Who is brack obama
Reformulated question: When was Barack obama born?
This reformulation is good since I am able to understand the question which I earlier was not able to understand without chat history."""
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 preknowledge. Third, utilizing the identified clues, reasoning, and input, furnish the pertinent answer for the question. Remember, you are required to use ONLY the provided preknowledge 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)


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 getQuestionAnswerOnTheBasisOfPreKnowledge(question, preKnowledge, systemPrompt):
    system_message = {"role": "system", "content": systemPrompt + preKnowledge}

    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):
    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 = getQuestionAnswerOnTheBasisOfPreKnowledge(
        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")

    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):
        urls = gr.Markdown()

    demo.load(
        lambda: ", ".join(
            random.sample(relevant_content, min(100, len(relevant_content)))
        ),
        None,
        urls,
    )

gr.close_all()
demo.launch()