import os
import shutil
import gradio as gr
import qdrant_client
from getpass import getpass


openai_api_key = os.getenv('OPENAI_API_KEY')

from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings

Settings.llm = OpenAI(model="gpt-3.5-turbo", temperature=0.4)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, StorageContext
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.core.memory import ChatMemoryBuffer


chat_engine = None
index = None
query_engine = None
memory = None
client = None
vector_store = None
storage_context = None


def process_upload(files):

    upload_dir = "uploaded_files"
    if not os.path.exists(upload_dir):
        os.makedirs(upload_dir)
    
    for file_path in files:
        file_name = os.path.basename(file_path)
        dest = os.path.join(upload_dir, file_name)
        if not os.path.exists(dest):
            shutil.copy(file_path, dest)
    
    documents = SimpleDirectoryReader(upload_dir).load_data()
    
    global client, vector_store, storage_context, index, query_engine, memory, chat_engine
    client = qdrant_client.QdrantClient(location=":memory:")
    
    vector_store = QdrantVectorStore(
        collection_name="paper",
        client=client,
        enable_hybrid=True,
        batch_size=20,
    )
    
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    
    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
    
    query_engine = index.as_query_engine(vector_store_query_mode="hybrid")
    
    memory = ChatMemoryBuffer.from_defaults(token_limit=3000)
    
    chat_engine = index.as_chat_engine(
        chat_mode="context",
        memory=memory,
        system_prompt=(
            "You are an AI assistant who answers the user questions"
        ),
    )
    
    return "Documents uploaded and index built successfully!"


def chat_with_ai(user_input, chat_history):
    global chat_engine
    if chat_engine is None:
        return chat_history, "Please upload documents first."
    
    response = chat_engine.chat(user_input)
    references = response.source_nodes
    ref, pages = [], []
    
    for node in references:
        file_name = node.metadata.get('file_name')
        if file_name and file_name not in ref:
            ref.append(file_name)
    
    complete_response = str(response) + "\n\n"
    if ref or pages:
        chat_history.append((user_input, complete_response))
    else:
        chat_history.append((user_input, str(response)))
    return chat_history, ""


def clear_history():
    return [], ""


def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# AI Assistant")
        
        with gr.Tab("Upload Documents"):
            gr.Markdown("Upload PDF, Excel, CSV, DOC/DOCX, or TXT files below:")
            # The file upload widget: we specify allowed file types.
            file_upload = gr.File(
                label="Upload Files",
                file_count="multiple",
                file_types=[".pdf", ".csv", ".txt", ".xlsx", ".xls", ".doc", ".docx"],
                type="filepath"
            )
            upload_status = gr.Textbox(label="Upload Status", interactive=False)
            upload_button = gr.Button("Process Upload")
            
            upload_button.click(process_upload, inputs=file_upload, outputs=upload_status)
        
        with gr.Tab("Chat"):
            chatbot = gr.Chatbot(label="Chatbot Assistant")
            user_input = gr.Textbox(
                placeholder="Ask a question...", label="Enter your question"
            )
            submit_button = gr.Button("Send")
            btn_clear = gr.Button("Restart")
            
            # A State to hold the chat history.
            chat_history = gr.State([])
            
            submit_button.click(chat_with_ai, inputs=[user_input, chat_history], outputs=[chatbot, user_input])
            user_input.submit(chat_with_ai, inputs=[user_input, chat_history], outputs=[chatbot, user_input])
            btn_clear.click(clear_history, outputs=[chatbot, user_input])
    
    return demo

# Launch the Gradio app.
gradio_interface().launch(debug=True)