import gradio as gr
from langchain.vectorstores import Qdrant
import qdrant_client
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import OpenAIEmbeddings
import dotenv
import os
from utils import template
import time

# Load environment variables and validate
#dotenv.load_dotenv()
QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
if not QDRANT_URL or not QDRANT_API_KEY:
    raise ValueError("QDRANT_URL and QDRANT_API_KEY must be set in the environment")

# Initialize the vector store
def initiliaze_vector_store():
    """
    Initialize and return the vector store.
    Only run this on launch.
    """
    embeddings = OpenAIEmbeddings()
    client = qdrant_client.QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
    vectorstore = Qdrant(client=client, collection_name="doc_datategy", embeddings=embeddings)
    return vectorstore

# Initialize the retriever
def initiliaze_retriever(vectorstore):
    """
    Initialize and return the retriever using the given vectorstore.
    """
    return vectorstore.as_retriever()

# Initialize the chatbot
def initiliaze_chatbot(template, model_name="gpt-3.5-turbo-1106", temperature=0):
    """
    Initialize and return the chatbot components: prompt and language model.
    """
    prompt = ChatPromptTemplate.from_template(template)
    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
    return prompt, llm

# Initialize the RAG chain
def initiliaze_RAG(retriever, prompt, llm):
    """
    Initialize and return the RAG chain.
    """
    context_function = {"context": retriever, "question": RunnablePassthrough()}
    rag_chain = context_function | prompt | llm | StrOutputParser()
    return rag_chain

# Launch Gradio app
vectorstore = initiliaze_vector_store()
retriever = initiliaze_retriever(vectorstore)

with gr.Blocks() as demo:

    welcome_text = """
    ### Chat with PapAI documentation  🚀

    **What's PapAI ? A revolutionary platform for Data Science projects and collaboration.** This application, powered by Gradio, offers a unique interactive experience by integrating the most advanced technologies in natural language processing 💬 and vector search 🔍 to answer your question related to the official papAI documentation (https://doc.datategy.net/)

    Through this interface, I use a combination of cutting-edge language models 🧠, such as GPT-3.5 Turbo, and a Qdrant vector storage system 🗃️ to provide accurate and relevant answers to a variety of Data Science queries. The core of PapAI lies in its ability to initialize and use RAG (Retrieval-Augmented Generation) chains 🔗, effectively combining information retrieval and text generation for an optimal user experience.
    
    Like all RAG applications, there is a possibility of encountering hallucinations, where the system might generate plausible but incorrect or nonsensical information."""

    # Add the Markdown widget to display the text
    gr.Markdown(welcome_text)

    # chatbot widget with chat history
    chatbot = gr.Chatbot(label="PapAI custom chatbot")
    msg = gr.Textbox(label="Prompt", value='What is the minimum configuration for running papAI ?', interactive=True)
    clear = gr.Button("Clear")
    template_user = gr.Textbox(label="Template", value=template, interactive=True)
    def change_template(template_user_str):
        prompt, llm = initiliaze_chatbot(template_user_str)
        return initiliaze_RAG(retriever, prompt, llm)

    def RAG_answer(query, chat_history, template_user_str):
        rag_chain = change_template(template_user_str)
        answer = rag_chain.invoke(query)
        chat_history.append((query, answer))
        time.sleep(0.5)  # Consider optimizing or dynamic handling
        return "", chat_history

    msg.submit(RAG_answer, [msg, chatbot, template_user], [msg, chatbot])

demo.queue()
demo.launch(share=False, debug=True)