import os
import torch
from langchain.chains import RetrievalQA
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.llms import HuggingFaceHub

# Check for GPU availability
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

# Global variables
conversation_retrieval_chain = None
chat_history = []
llm_hub = None
embeddings = None

def init_llm():
    global llm_hub, embeddings

    # Ensure API key is set in Hugging Face Spaces
    hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
    if not hf_token:
        raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set in environment variables.")

    model_id = "tiiuae/falcon-7b-instruct"
    llm_hub = HuggingFaceHub(repo_id=model_id, model_kwargs={"temperature": 0.1, "max_new_tokens": 600, "max_length": 600})

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": DEVICE}
    )

def process_document(document_path):
    global conversation_retrieval_chain

    # Ensure LLM and embeddings are initialized
    if not llm_hub or not embeddings:
        init_llm()

    loader = PyPDFLoader(document_path)
    documents = loader.load()
    
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
    texts = text_splitter.split_documents(documents)
    
    db = Chroma.from_documents(texts, embedding=embeddings, persist_directory="./chroma_db")

    conversation_retrieval_chain = RetrievalQA.from_chain_type(
        llm=llm_hub,
        chain_type="stuff",
        retriever=db.as_retriever(search_type="mmr", search_kwargs={'k': 6, 'lambda_mult': 0.25}),
        return_source_documents=False
    )

def process_prompt(prompt):
    global conversation_retrieval_chain, chat_history

    if not conversation_retrieval_chain:
        return "No document has been processed yet. Please upload a PDF first."

    output = conversation_retrieval_chain({"query": prompt, "chat_history": chat_history})
    answer = output["answer"]
    
    chat_history.append((prompt, answer))
    
    return answer