# pdf_chat.py - a small RAG chatbot: ask questions about one PDF.
# Run: python pdf_chat.py leave_policy.pdf
import sys
from pypdf import PdfReader
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from google import genai

PDF_FILE = sys.argv[1] if len(sys.argv) > 1 else "leave_policy.pdf"
MODEL = "gemini-3.8-flash"   # check the current model name in Google's Gemini API docs

# Step 1. LOAD: read the text of every page
reader = PdfReader(PDF_FILE)
pages = [(n + 1, page.extract_text() or "") for n, page in enumerate(reader.pages)]

# Step 2. CHUNK: cut each page into pieces of about 80 words (20 words overlap)
chunks = []
for page_no, text in pages:
    words = text.split()
    for start in range(0, len(words), 60):
        piece = " ".join(words[start:start + 80])
        if piece:
            chunks.append((page_no, piece))
print(f"Loaded {len(pages)} pages and made {len(chunks)} chunks from {PDF_FILE}")

# Step 3. INDEX: turn every chunk into numbers so we can search it
vectorizer = TfidfVectorizer(stop_words="english")
matrix = vectorizer.fit_transform([text for _, text in chunks])

# Step 4. RETRIEVE: find the chunks that best match the question
def retrieve(question, top_k=3):
    scores = cosine_similarity(vectorizer.transform([question]), matrix)[0]
    best = scores.argsort()[::-1][:top_k]
    return [(chunks[i][0], chunks[i][1]) for i in best if scores[i] > 0]

# Step 5. GENERATE: give only those chunks to the model, with strict rules
client = genai.Client()   # reads your key from the GEMINI_API_KEY environment variable

def answer(question):
    found = retrieve(question)
    if not found:
        return "I could not find this in the document."
    context = "\n\n".join(f"[page {p}] {text}" for p, text in found)
    prompt = (
        "Answer the question using ONLY the context below.\n"
        'If the answer is not in the context, reply: "I could not find this in the document."\n'
        "Keep it short and mention the page number you used.\n\n"
        f"Context:\n{context}\n\nQuestion: {question}"
    )
    result = client.interactions.create(model=MODEL, input=prompt)
    return result.output_text

while True:
    q = input("\nAsk a question (or type exit): ").strip()
    if q.lower() == "exit":
        break
    if q.lower().startswith("show "):            # debug: see what was retrieved
        for p, text in retrieve(q[5:]):
            print(f"--- page {p}: {text[:120]}...")
        continue
    print(answer(q))