import os
import datetime
import time
import gradio as gr

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
# Replacing vector components with a statistical BM25 (vectorless) retriever
from langchain_community.retrievers import BM25Retriever
from langchain_community.llms import LlamaCpp
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

# -------------------------------------------------------------------------
# 1. Global Setup (LLM Loaded Once - Embeddings Removed)
# -------------------------------------------------------------------------
print("Initializing Local LLM Runtime (Vectorless Mode)...")

MODEL_PATH = "model/gemma-4-E2B-it-Q8_0.gguf"  # Ensure this path is correct
llm = LlamaCpp(
    model_path=MODEL_PATH,
    n_ctx=4096,
    temperature=0.1,
    max_tokens=512, 
    verbose=False
)

# Build the rigid system prompt structure
system_prompt = (
    "You are a strict Document Retrieval analysis assistant.\n"
    "Answer the user's question using ONLY the following pieces of retrieved financial context. "
    "If you do not know the answer or if it is not explicitly stated in the context, "
    "state clearly that the book does not provide this information. Do not make up facts.\n\n"
    "Context:\n{context}"
)
prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}"),
])
question_answer_chain = create_stuff_documents_chain(llm, prompt)

# -------------------------------------------------------------------------
# 2. Core Processing Functions
# -------------------------------------------------------------------------
def process_pdf(pdf_file):
    """Loads, splits, and indexes the PDF using BM25 string-matching instead of vectors."""
    if pdf_file is None:
        return "No file uploaded.", None, ""

    start_time = time.time()
    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    log_output = f"[{current_time}] 📁 Loading PDF: {os.path.basename(pdf_file.name)}...\n"
    
    try:
        # Load and Split
        loader = PyPDFLoader(pdf_file.name)
        docs = loader.load()
        
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        final_chunks = text_splitter.split_documents(docs)
        log_output += f" -> Successfully generated {len(final_chunks)} semantic chunks.\n"
        
        # Build Vectorless BM25 Text Index
        log_output += f" -> Indexing text chunks into local BM25 Search Engine...\n"
        bm25_retriever = BM25Retriever.from_documents(final_chunks)
        # Configure to return top 3 keyword-matched document chunks
        bm25_retriever.k = 3 
        
        elapsed = time.time() - start_time
        log_output += f"✅ Success! Vectorless processing finished in {elapsed:.2f} seconds."
        
        return log_output, bm25_retriever, "PDF text indexed! You can now ask questions."
        
    except Exception as e:
        return f"❌ Error processing PDF: {str(e)}", None, "Processing failed."

def answer_query(query, bm25_retriever, chat_history):
    """Retrieves context using BM25 and runs inference against the local GGUF model."""
    current_time = datetime.datetime.now().strftime("%H:%M:%S")
    
    if bm25_retriever is None:
        chat_history.append((query, "⚠️ Please upload and process a financial PDF first!"))
        return "", chat_history

    if not query.strip():
        return "", chat_history

    # Dynamically bind the pre-computed BM25 engine into the RAG chain
    rag_chain = create_retrieval_chain(bm25_retriever, question_answer_chain)
    
    # Run Inference
    response = rag_chain.invoke({"input": query})
    answer = response['answer']
    
    # Format answer with timestamp marker
    formatted_answer = f"[{current_time}]\n{answer}"
    
    chat_history.append((query, formatted_answer))
    return "", chat_history

# -------------------------------------------------------------------------
# 3. Gradio UI Architecture
# -------------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # State now holds the BM25Retriever object instead of a FAISS DB
    retriever_state = gr.State(None)
    
    gr.Markdown(
        """
        # 🏦 Secure Document Retrieve Vectorless RAG Engine
        Analyze private documents completely offline using high-performance local AI and keyword-based retrieval.
        """
    )
    
    with gr.Row():
        # Left Panel: Document Ingestion and Pipeline Metrics
        with gr.Column(scale=1):
            gr.Markdown("### 1. Document Ingestion")
            file_input = gr.File(label="Upload PDF/ Book", file_types=[".pdf"])
            process_btn = gr.Button("Build Text Keyword Base", variant="primary")
            
            status_text = gr.Textbox(label="Engine Status", interactive=False, value="Awaiting document...")
            log_monitor = gr.TextArea(label="Pipeline Terminal Logs", interactive=False, lines=8)
            
        # Right Panel: Interactive Chat Interface
        with gr.Column(scale=2):
            gr.Markdown("### 2. Conversational Analysis")
            chatbot = gr.Chatbot(label="Document Retrieval Chat", height=450)
            
            with gr.Row():
                query_input = gr.Textbox(
                    label="Ask a document retrieval question...",
                    placeholder="e.g., What is the primary difference between working capital and capital budgeting?",
                    scale=4
                )
                submit_btn = gr.Button("Submit Query", variant="secondary", scale=1)

    # --- UI Event Wire-up ---
    
    # File processing sequence
    process_btn.click(
        fn=process_pdf,
        inputs=[file_input],
        outputs=[log_monitor, retriever_state, status_text]
    )
    
    # Query execution sequences (Triggers on Click or hitting Enter)
    submit_btn.click(
        fn=answer_query,
        inputs=[query_input, retriever_state, chatbot],
        outputs=[query_input, chatbot]
    )
    query_input.submit(
        fn=answer_query,
        inputs=[query_input, retriever_state, chatbot],
        outputs=[query_input, chatbot]
    )

# Launch local server
if __name__ == "__main__":
    demo.launch(server_name="127.0.0.1", server_port=7860, share=False)