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
import base64

# Load environment variables
load_dotenv()

# Default models
selected_llm_model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
selected_embed_model_name = "BAAI/bge-small-en-v1.5"
vector_index = None

# Initialize the parser
parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
file_extractor = {
    '.pdf': parser, '.docx': parser, '.doc': parser, '.txt': parser,
    '.csv': parser, '.xlsx': parser, '.pptx': parser, '.html': parser,
    '.jpg': parser, '.jpeg': parser, '.png': parser, '.webp': parser, '.svg': parser
}

# 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()
        embed_model = HuggingFaceEmbedding(model_name=selected_embed_model_name)
        vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
        print(f"Parsing done for {file_path}")
        filename = os.path.basename(file_path)
        return f"File upload status: Ready to write ({filename})"
    except Exception as e:
        return f"An error occurred: {e}"

# Respond function
def respond(message, history):
    try:
        llm = HuggingFaceInferenceAPI(
            model_name=selected_llm_model_name,
            contextWindow=8192, maxTokens=1024, temperature=0.3, topP=0.9,
            frequencyPenalty=0.5, presencePenalty=0.5, token=os.getenv("TOKEN")
        )
        query_engine = vector_index.as_query_engine(llm=llm)
        bot_message = query_engine.query(message)
        print(f"\n{datetime.now()}:{selected_llm_model_name}:: {message} --> {str(bot_message)}\n")
        return f"{selected_llm_model_name}:\n{str(bot_message)}"
    except Exception as e:
        if str(e) == "'NoneType' object has no attribute 'as_query_engine'":
            return "Please upload a file."
        return f"An error occurred: {e}"

# UI Setup
with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]), css='footer {visibility: hidden}') as demo:
    gr.Markdown("# Document RAG")
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(file_count="single", type='filepath', label="Upload Document")
            btn = gr.Button("Submit", variant='primary')
            clear = gr.ClearButton()
            output = gr.Text(label='File Upload Status')
        with gr.Column(scale=3):
            gr.ChatInterface(
                fn=respond,
                chatbot=gr.Chatbot(height=500),
                theme="light",
                show_progress='full',
                textbox=gr.Textbox(placeholder="Ask me questions on the uploaded document!", container=False)
            )
    btn.click(fn=load_files, inputs=[file_input], outputs=output)
    clear.click(lambda: [None] * 2, outputs=[file_input, output])

# Launch the demo
if __name__ == "__main__":
    demo.launch()