import os
import gradio as gr
from pinecone import Pinecone
from langchain_chroma import Chroma
from langchain_core.prompts import PromptTemplate
from langchain_pinecone import PineconeVectorStore
from langchain_community.vectorstores import FAISS
from langchain_community.vectorstores import LanceDB
from google.api_core.exceptions import ResourceExhausted
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_google_genai import GoogleGenerativeAIEmbeddings, GoogleGenerativeAI

embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
gemini = GoogleGenerativeAI(model="models/gemini-2.0-flash", temperature=0.0, top_k=1, top_p=0.0)

prompt_template = """
    
    Context:\n {context}?\n
    Question: \n{question}\n

    Answer:
    """

prompt = PromptTemplate(template = prompt_template, input_variables = ["context", "question"])

chain = prompt | gemini.with_retry(retry_if_exception_type=(ResourceExhausted,))

index_name = "langchain-test-index"

def extract_text_from_pdf(pdf_path):
    raw_documents = []
    for path in pdf_path:
        raw_documents.extend(PyPDFLoader(path).load())
    return raw_documents

def chunk_text(raw_documents, chunk_size, chunk_overlap):
    text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    documents = text_splitter.split_documents(raw_documents)
    return documents

def delete_pinecone():
    pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])

    index = pc.Index(host="https://langchain-test-index-la2n80y.svc.aped-4627-b74a.pinecone.io")

    if index.describe_index_stats()['total_vector_count'] > 0:
        index.delete(delete_all=True)

def store_chroma_db(documents):
    chroma_db = Chroma.from_documents(documents, embeddings, persist_directory="./chroma_db")

def store_faiss_db(documents):
    faiss_db = FAISS.from_documents(documents, embeddings)
    faiss_db.save_local("./faiss_db")

def store_lance_db(documents):
    lance_db = LanceDB.from_documents(documents, embeddings, uri="./lance_db")

def store_pinecone_db(documents):
    pinecone_db = PineconeVectorStore.from_documents(documents, index_name=index_name,
                                                     embedding=embeddings)
    
def load_chroma_db():
    chroma_db = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
    return chroma_db

def load_faiss_db():
    faiss_db = FAISS.load_local("./faiss_db", embeddings, allow_dangerous_deserialization=True)
    return faiss_db

def load_lance_db():
    lance_db = LanceDB(embedding=embeddings, uri="./lance_db")
    return lance_db

def connect_pinecone_db():
    pinecone_db = PineconeVectorStore(index_name=index_name, embedding=embeddings)
    return pinecone_db

def invoke_chain(db, query):
    docs = db.similarity_search(query)
    answer = chain.invoke({"context":docs, "question": query}, return_only_outputs=True)
    return answer
    
def store_embeddings(pdf_path, chunk_size, chunk_overlap):
    raw_documents = extract_text_from_pdf(pdf_path)

    documents = chunk_text(raw_documents, chunk_size, chunk_overlap)

    delete_pinecone()

    store_chroma_db(documents)
    store_faiss_db(documents)
    store_lance_db(documents)
    store_pinecone_db(documents)
    
    return "All embeddings are stored in vector database"

title = "PDF Chat"
description = "A simple Gradio interface to query PDFs and compare vector database"
examples = [[["data/amazon-10-k-2024.pdf"], 1000, 100],
            [["data/goog-10-k-2023.pdf"], 1000, 100]]

def inference(query):
    chroma_db = load_chroma_db()
    chroma_answer = invoke_chain(chroma_db, query)

    faiss_db = load_faiss_db()
    faiss_answer = invoke_chain(faiss_db, query)

    lance_db = load_lance_db()
    lance_answer = invoke_chain(lance_db, query)

    pinecone_db = connect_pinecone_db()
    pinecoce_answer = invoke_chain(pinecone_db, query)

    return chroma_answer, faiss_answer, lance_answer, pinecoce_answer


with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown(f"# {title}\n{description}")
    with gr.Row():
        with gr.Column():
            pdf = gr.File(label="Input PDFs", file_count="multiple", file_types=[".pdf"])
            chunk_size = gr.Slider(0, 2000, 1000, 100, label="Size of Chunk")
            chunk_overlap = gr.Slider(0, 1000, 100, 100, label="Size of Chunk Overlap")
            with gr.Row():
                clear_btn = gr.ClearButton(components=[pdf, chunk_size, chunk_overlap])
                submit_btn = gr.Button("Store Embeddings", variant='primary')
        with gr.Column():
            message = gr.Textbox(label="Status", type="text")
            chroma_out = gr.Textbox(label="ChromaDB Response", type="text")
            faiss_out = gr.Textbox(label="FaissDB Response", type="text")
            lance_out = gr.Textbox(label="LanceDB Response", type="text")
            pinecone_out = gr.Textbox(label="PineconeDB Response", type="text")

    submit_btn.click(store_embeddings, inputs=[pdf, chunk_size, chunk_overlap], outputs=message)

    with gr.Row():
        with gr.Column():
            text = gr.Textbox(label="Question", type="text")
            with gr.Row():
                chat_clear_btn = gr.ClearButton(components=[text])
                chat_submit_btn = gr.Button("Submit", variant='primary')

    chat_submit_btn.click(inference, inputs=[text], outputs=[chroma_out, faiss_out, lance_out,
                                                             pinecone_out])

    examples_obj = gr.Examples(examples=examples, inputs=[pdf, chunk_size, chunk_overlap])

demo.launch()