app.py

import gradio as gr
import os
import shutil
from src.config import PDF_PATH, HF_API_TOKEN, HF_MODEL_ID, DATA_DIR
from src.ingestion import ingest_file
from src.vectorstore import create_vectorstore, load_vectorstore
from src.agent import build_langgraph_agent
from langchain_core.messages import HumanMessage

# Global variables to store state
vectorstore = None
agent_executor = None
current_hf_token = None
current_hf_model = None

# Ensure data directory exists for uploads and FAISS index (important for HF Spaces).
os.makedirs(DATA_DIR, exist_ok=True)


def _get_uploaded_path(uploaded_file):
    """
    Normalize Gradio's uploaded file into a filesystem path.
    Handles filepath strings, temporary file objects, and dict payloads.
    """
    if uploaded_file is None:
        return None

    if isinstance(uploaded_file, (str, os.PathLike)):
        return str(uploaded_file)

    if isinstance(uploaded_file, dict):
        return uploaded_file.get("name") or uploaded_file.get("path")

    if hasattr(uploaded_file, "name"):
        return uploaded_file.name

    return None


def initialize_system(hf_token, hf_model, uploaded_file):
    """
    Initializes the RAG pipeline and Agent.
    """
    global vectorstore, agent_executor, current_hf_token, current_hf_model
    
    hf_token = (hf_token or HF_API_TOKEN or "").strip()
    hf_model = (hf_model or HF_MODEL_ID).strip()
    uploaded_path = _get_uploaded_path(uploaded_file)
    
    if uploaded_file is not None and uploaded_path is None:
        return "Could not read the uploaded file. Please try uploading again."

    if uploaded_path is None and not os.path.exists(PDF_PATH):
        return "Please upload a PDF file."

    try:
        # 0. Handle File Upload
        if uploaded_path is not None:
            # Gradio passes a temporary file path or a file object depending on version/config.
            # Usually it's a named temp file path in recent versions.
            # We copy it to our data directory.
            if not os.path.exists(os.path.dirname(PDF_PATH)):
                os.makedirs(os.path.dirname(PDF_PATH))
                
            # uploaded_file is a file path in recent Gradio versions
            shutil.copy(uploaded_path, PDF_PATH)
            print(f"File saved to {PDF_PATH}")
            
            # Force re-ingestion since we have a new file
            print("Ingesting PDF...")
            chunks = ingest_file(str(PDF_PATH))
            vectorstore = create_vectorstore(chunks)
        
        # 1. Load or Create Vector Store (if not already created above)
        if vectorstore is None:
            vectorstore = load_vectorstore()
            if vectorstore is None:
                # This case should be covered by the upload logic, but just in case
                if os.path.exists(PDF_PATH):
                    print("Ingesting PDF...")
                    chunks = ingest_file(str(PDF_PATH))
                    vectorstore = create_vectorstore(chunks)
                else:
                    return "Source PDF not found. Please upload a file."
        
        # 2. Create Agent (LangGraph)
        agent_executor = build_langgraph_agent(vectorstore, hf_api_token=hf_token, hf_model_id=hf_model)
        current_hf_token = hf_token
        current_hf_model = hf_model
        mode = "Hugging Face router" if hf_token else "local transformers (no HF token provided)"
        
        return f"System Initialized Successfully using {mode}. You can now start chatting."
    except Exception as e:
        import traceback
        traceback.print_exc()
        return f"Initialization Failed: {str(e)}"

def chat(message, history, hf_token, hf_model, uploaded_file):
    """
    Chat function for Gradio.
    """
    global agent_executor, current_hf_token, current_hf_model
    
    # Gradio can pass None for history on the first turn.
    history = history or []
    if not message:
        return "Please enter a message to start chatting."
        
    hf_token = (hf_token or HF_API_TOKEN or "").strip()
    hf_model = (hf_model or HF_MODEL_ID).strip()
    
    # Check if API key has changed or agent is not initialized
    if agent_executor is None or hf_token != current_hf_token or hf_model != current_hf_model:
        init_msg = initialize_system(hf_token, hf_model, uploaded_file)
        if "Failed" in init_msg or "Please" in init_msg:
            return init_msg
    
    # Run the agent
    try:
        # Convert history to LangChain format if needed, but LangGraph handles state.
        # We pass the full history + new message to the agent if we were managing state manually,
        # but here we'll just pass the new message and let the graph handle it if we were persistent.
        # For a simple chat interface without persistence, we pass the conversation history.
        
        messages = []
        for h in history:
            messages.append(HumanMessage(content=h[0]))
            # We would need AI message here too, but Gradio history is [user, bot].
            # For simplicity in this demo, we'll just send the current message or a limited context.
            # Let's send the current message. To support history, we'd need to map Gradio history to LangChain messages.
        
        # Better approach for this demo: Just send the current message. 
        # The agent is stateless between calls in this simple implementation unless we use checkpointers.
        
        response = agent_executor.invoke({"messages": [HumanMessage(content=message)]})
        return response["messages"][-1].content
    except Exception as e:
        import traceback
        traceback.print_exc()
        hint = (
            " If you used the Hugging Face router, verify the token/model. "
            "Otherwise, try re-initializing to refresh the vector store."
        )
        return f"Error while generating a reply: {str(e)}{hint}"

# Gradio UI
with gr.Blocks(title="RAG Chatbot (LangGraph + HF)") as demo:
    gr.Markdown("# RAG-Based Chatbot (LangGraph + Hugging Face)")
    gr.Markdown(
        "Upload a PDF, build a vector store, retrieve context, and answer with either the Hugging Face router "
        "(when a token + router model is provided) or a local fallback model."
    )
    
    with gr.Row():
        api_key_input = gr.Textbox(
            label="Hugging Face API Token (optional)", 
            type="password", 
            placeholder="hf_...",
            value=os.getenv("HUGGINGFACEHUB_API_TOKEN", "")
        )
        model_input = gr.Textbox(
            label="Model ID",
            placeholder="e.g. meta-llama/Meta-Llama-3-8B-Instruct",
            value=os.getenv("HF_MODEL_ID", HF_MODEL_ID),
        )
        file_input = gr.File(label="Upload PDF", file_types=[".pdf"], type="filepath")
        init_btn = gr.Button("Initialize System")
    
    status_output = gr.Textbox(label="Status", interactive=False)
    
    chatbot = gr.ChatInterface(
        fn=chat, 
        additional_inputs=[api_key_input, model_input, file_input]
    )

    init_btn.click(initialize_system, inputs=[api_key_input, model_input, file_input], outputs=[status_output])

if __name__ == "__main__":
    # Use local launch by default; share links can fail without network access.
    demo.queue().launch(share=False)