import os
import time
from huggingface_hub import login
import gradio as gr
import traceback
from pypdf import PdfReader
import uuid

# --- 1. Authentication ---
# This uses the secret named 'token' you created in Space Settings
# HF_TOKEN = os.environ.get("token")
# if HF_TOKEN:
#     login(token=HF_TOKEN)

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_groq import ChatGroq
from langchain_classic.chains import ConversationalRetrievalChain

# --- Configuration & Constraints ---
MAX_PDFS = 2
MAX_PAGES = 3
EMBEDDING_MODEL = "BAAI/bge-small-en-v1.5"
GROQ_MODEL = "llama-3.3-70b-versatile"

# Initialize embeddings locally using the token for access
embeddings = HuggingFaceEmbeddings(
    model_name=EMBEDDING_MODEL,
    model_kwargs={'device': 'cpu'}
)


def is_valid_pdf(file):
    try:
        # We only read the header/metadata, not the whole file
        reader = PdfReader(file.name)
        # Accessing the length of pages triggers a basic structure check
        if len(reader.pages) > 0:
            return True
        return False
    except Exception:
        return False

def process_pdfs(files):
    if not files:
        return None, "❌ No files provided."
    
    if len(files) > MAX_PDFS:
        return None, f"❌ Error: Max {MAX_PDFS} PDFs allowed."

    for file in files:
        if not is_valid_pdf(file):
            return None, "❌ PDF not valid."
        

    all_docs = []
    for file in files:
        loader = PyPDFLoader(file.name)
        pages = loader.load()
        if len(pages) > MAX_PAGES:
            return None, f"❌ Error: '{os.path.basename(file.name)}' exceeds {MAX_PAGES} pages."
        all_docs.extend(pages)

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    splits = text_splitter.split_documents(all_docs)
    vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
    
    return vectorstore, "✅ PDFs processed successfully!"

def chat_with_pdf(message, history, vectorstore):
    # 1. Ensure history is a list (Gradio 6 initialization)
    if history is None:
        history = []

    # 2. Handle the 'None' case for vectorstore (before processing PDFs)
    if vectorstore is None:
        # Standard dictionary format for Gradio 6
        history.append({"role": "user", "content": str(message)})
        history.append({"role": "assistant", "content": "Please upload and process PDFs first."})
        return history, ""

    llm = ChatGroq(
        temperature=0,
        model_name=GROQ_MODEL,
        groq_api_key=os.environ.get("GROQ_API_KEY") 
    )

    # 3. Correctly format history for LangChain (List of Tuples)
    formatted_history = []
    
    user_msg = None
    
    for msg in history:
        role = msg.get("role")
        content = msg.get("content", "")
    
        if role == "user":
            user_msg = content
    
        elif role == "assistant" and user_msg is not None:
            formatted_history.append((user_msg, content))
            user_msg = None

    qa_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
    )

    try:
        # 4. Invoke the chain
        # result = qa_chain.invoke({"question": str(message), "chat_history": formatted_history})
        result = qa_chain.invoke({
            "question": str(message),
            "chat_history": [
                (str(u), str(a)) for u, a in formatted_history
            ]
        })
        answer = result['answer']
    except Exception as e:
        print("❌ FULL EXCEPTION TRACEBACK:")
        print(traceback.format_exc())
    
        error_str = str(e).lower()
        if "429" in error_str:
            answer = "⚠️ [Rate Limit] System is busy. Please wait a moment."
        else:
            answer = f"⚠️ System Error: {str(e)}"

    # 5. Append the new interaction to history
    history.append({"role": "user", "content": str(message)})
    history.append({"role": "assistant", "content": str(answer)})
    
    return history, ""

# --- UI Setup ---
# Note: theme removed from Blocks() per Gradio 6.0 warning
with gr.Blocks() as demo:
    vector_db = gr.State(None)
    
    gr.Markdown(f"# 🚀 PDF RAG Demo\nMax {MAX_PDFS} PDFs | Max {MAX_PAGES} pages each")

    with gr.Row():
        with gr.Column(scale=1):
            file_uploader = gr.File(label="Upload PDFs", file_count="multiple")
            process_btn = gr.Button("Build Knowledge Base", variant="primary")
            status_msg = gr.Textbox(label="Status", interactive=False)

        with gr.Column(scale=2):
            # 'type' argument removed as it is now default/standard
            chatbot = gr.Chatbot(label="Chat History", height=450)
            with gr.Row():
                question_input = gr.Textbox(label="Your Question", placeholder="Ask away...", scale=4)
                submit_btn = gr.Button("Send", scale=1)
            clear_btn = gr.ClearButton([question_input, chatbot])

    process_btn.click(
        process_pdfs, 
        inputs=[file_uploader], 
        outputs=[vector_db, status_msg], 
        api_name="upload_and_process"
    )
    
    submit_btn.click(
        chat_with_pdf, 
        inputs=[question_input, chatbot, vector_db], 
        outputs=[chatbot, question_input], 
        api_name="ask_pdf"
    )
    
    question_input.submit(
        chat_with_pdf, 
        inputs=[question_input, chatbot, vector_db], 
        outputs=[chatbot, question_input]
    )

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