import os
import gradio as gr
from dotenv import load_dotenv
from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler, CBEventType
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.llms.openai import OpenAI
from llama_index.llms.groq import Groq
from llama_index.core.base.embeddings.base import similarity
from llama_index.llms.ollama import Ollama
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core import StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import load_index_from_storage


load_dotenv()

# set up LLM
GROQ_API_KEY = os.getenv('GROQ_API_KEY')
llm = Groq(model="llama3-70b-8192")
Settings.llm = llm

# set up callback manager
llama_debug = LlamaDebugHandler(print_trace_on_end=True)
callback_manager = CallbackManager([llama_debug])
Settings.callback_manager = callback_manager

# converting documents into embeddings and indexing
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.embed_model = embed_model

# create splitter
splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=20)
Settings.transformations = [splitter]

if os.path.exists("./vectordb"):
    storage_context = StorageContext.from_defaults(persist_dir="./vectordb")
    index = load_index_from_storage(storage_context)
else:
    filename_fn = lambda filename: {"file_name": filename}
    required_exts = [".pdf",".docx"]
    reader = SimpleDirectoryReader(
        input_dir="./data",
        required_exts=required_exts,
        recursive=True,
        file_metadata=filename_fn
    )
    documents = reader.load_data()
    for doc in documents:
        doc.text = str(doc.metadata) +' '+ doc.text 
    print("index creating with `%d` documents", len(documents))
    index = VectorStoreIndex.from_documents(documents, embed_model=embed_model, text_splitter=splitter)

    index.storage_context.persist(persist_dir="./vectordb")

# set up query engine
query_engine = index.as_query_engine(
    similarity_top_k=5,
    #node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.7)],
    verbose=True,
)

def retreive(question):
    qns_w_source = "Answer the following question: " + question + " Followed by providing the page and file name of the source document as well, thank you!"
    streaming_response = query_engine.query(qns_w_source)
    #sources = streaming_response.get_formatted_sources(length=5000)
    return str(streaming_response) # + "\n" + str(sources)

demo = gr.Interface(fn=retreive, inputs="textbox", outputs="textbox")

if __name__ == "__main__":
    demo.launch()