import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain_groq import ChatGroq
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough


# Initialize the FAISS vector store
vector_store = None

# Function to handle PDF upload and indexing
def index_pdf(pdf):
    global vector_store
    
    # Load the PDF
    loader = PyPDFLoader(pdf.name)
    documents = loader.load()

    # Split the documents into chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    texts = text_splitter.split_documents(documents)

    # Embed the chunks 
    embeddings = HuggingFaceEmbeddings(model_name="bert-base-uncased", encode_kwargs={"normalize_embeddings": True})

    # Store the embeddings in the vector store
    vector_store = FAISS.from_documents(texts, embeddings)

    return "PDF indexed successfully!"

# Function to handle chatbot queries
def chatbot_query(query):
    if vector_store is None:
        return "Please upload and index a PDF first."
    
    # Create a retrieval-based QA chain
    retriever = vector_store.as_retriever()
    qa_chain = RetrievalQA(llm=OpenAI(), retriever=retriever)
    
    # Get the response from the QA chain
    response = qa_chain.run(query)
    
    return response


def generate_response(query, history, model, temperature, max_tokens, top_p, seed):
    response = query + "\n"
    response = response + model + "\n"

    return response



additional_inputs = [
    gr.Dropdown(choices=["llama-3.1-70b-versatile", "llama-3.1-8b-instant", "llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma2-9b-it", "gemma-7b-it"], value="llama-3.1-70b-versatile", label="Model"),
    gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Temperature", info="Controls diversity of the generated text. Lower is more deterministic, higher is more creative."),
    gr.Slider(minimum=1, maximum=8000, step=1, value=8000, label="Max Tokens", info="The maximum number of tokens that the model can process in a single response.<br>Maximums: 8k for gemma 7b it, gemma2 9b it, llama 7b & 70b, 32k for mixtral 8x7b, 132k for llama 3.1."),
    gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Top P", info="A method of text generation where a model will only consider the most probable next tokens that make up the probability p."),
    gr.Number(precision=0, value=0, label="Seed", info="A starting point to initiate generation, use 0 for random")
]


# Create the Gradio interface
with gr.Blocks(theme="Nymbo/Alyx_Theme") as demo:
    with gr.Tab("Indexing"):
        pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
        index_button = gr.Button("Index PDF")
        index_output = gr.Textbox(label="Indexing Status")
        index_button.click(index_pdf, inputs=pdf_input, outputs=index_output)
    
    with gr.Tab("Chatbot"):
        # query_input = gr.Textbox(label="Enter your question")
        # query_button = gr.Button("Submit")
        # query_output = gr.Textbox(label="Response")
        
        # query_button.click(chatbot_query, inputs=query_input, outputs=query_output)

        gr.ChatInterface(
            fn=generate_response, 
            chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
            additional_inputs=additional_inputs,
        )       

# Launch the Gradio app
demo.launch()