from datetime import datetime
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_parse import LlamaParse
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
import os
from dotenv import load_dotenv
import gradio as gr

# Load environment variables
load_dotenv()

# Initialize the LLM and parser
llm = HuggingFaceInferenceAPI(
    # model_name="meta-llama/Meta-Llama-3-8B-Instruct",
    model_name="mistralai/Mixtral-8x7B-Instruct-v0.1",
    token=os.getenv("TOKEN"),
    model_config={'protected_namespaces': ()}
)

parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
file_extractor = {'.pdf': parser, '.docx': parser, '.doc': parser}

# Embedding model and index initialization (to be populated by uploaded files)
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

# Global variable to store documents loaded from user-uploaded files
vector_index = None

# File processing function
def load_files(file_path: str):
    try:
        global vector_index
        document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
        vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
        print(f"parsing done {file_path}")
        filename = os.path.basename(file_path)
        return f"Ready to give response on give {filename}"
    except Exception as e:
        return f"An error occurred {e}"

def respond(message, history):
    try:
        query_engine = vector_index.as_query_engine(llm=llm)
        bot_message = query_engine.query(message)
        # output = ""
        # for chr in bot_message:
        #     output += chr
        #     yield output
        print(f"\n{datetime.now()}:: {message} --> {str(bot_message)}\n")
        return str(bot_message)
    except Exception as e:
        if e == "'NoneType' object has no attribute 'as_query_engine'":
            return "upload file"
        return f"an error occurred {e}"

# UI Setup
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(file_count="single", type='filepath')
            with gr.Row():
                clear = gr.ClearButton()
                btn = gr.Button("Submit", variant='primary')
            output = gr.Text(label='Vector Index')
        with gr.Column(scale=3):
             gr.ChatInterface(fn=respond,
                              chatbot=gr.Chatbot(height=500),
                              textbox=gr.Textbox(placeholder="Ask me questions on given document!", container=False, scale=7),
                              # examples=["summarize the document"]
                             )

    # Action on button click to process file and load into index
    btn.click(fn=load_files, inputs=file_input, outputs=output)
    clear.click(lambda: [None]*2, outputs=[file_input, output])


# Launch the demo with public link option
if __name__ == "__main__":
    demo.launch()