"""
RAG Document Q&A
----------------
Upload documents (PDF / .txt) or use the built-in sample job postings, then ask
questions. Answers are generated by an LLM but grounded ONLY in the retrieved
passages, with inline [n] citations back to the source.

Pipeline:  ingest -> chunk -> embed (sentence-transformers) -> index (FAISS)
           -> retrieve (cosine top-k) -> generate (Groq, OpenAI-compatible API)

Everything except the final generation runs locally on the free Spaces CPU.
Generation is a single API call to Groq's free tier (no credit card).
"""

import os
import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer
import faiss
from pypdf import PdfReader
from openai import OpenAI

# ----------------------------- Config --------------------------------------
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"  # ~80 MB, runs on CPU
GROQ_MODEL = "openai/gpt-oss-20b"  # current Groq production model.
# Model IDs change. List active ones any time at:
# https://console.groq.com/docs/models  (or GET /openai/v1/models)
CHUNK_SIZE = 200       # words per chunk
CHUNK_OVERLAP = 40     # words shared between neighbouring chunks
TOP_K = 4              # passages retrieved per question

# Load the embedding model once at startup (reused for every request).
embedder = SentenceTransformer(EMBED_MODEL)


def get_llm_client():
    """Return a Groq client via the OpenAI-compatible API, or None if no key."""
    key = os.environ.get("GROQ_API_KEY")
    if not key:
        return None
    return OpenAI(api_key=key, base_url="https://api.groq.com/openai/v1")


# ------------------- Sample corpus (so the demo works instantly) -----------
# Three short, realistic working-student postings with overlapping and
# distinct skills, so the example questions return interesting answers.
SAMPLE_DOCS = {
    "werkstudent_data_science_munich.txt": (
        "Working Student Data Science (m/w/d), Munich. You will support our team "
        "in building data pipelines and prototyping machine learning models for "
        "product analytics. Requirements: currently enrolled in computer science, "
        "data science, statistics or a related field. Strong Python skills and "
        "solid knowledge of SQL. Experience with pandas and scikit-learn. "
        "Familiarity with a cloud platform (AWS preferred) is a plus. English is "
        "our working language; German is advantageous but not required. "
        "16 hours per week, hybrid."
    ),
    "ml_engineer_working_student_berlin.txt": (
        "Machine Learning Working Student (m/f/d), Berlin. Join us to bring models "
        "into production. You will work on deep learning models using PyTorch, help "
        "deploy them as APIs, and experiment with large language models and "
        "retrieval-augmented generation. Requirements: completed coursework in "
        "machine learning or neural networks, good Python, hands-on experience with "
        "PyTorch or TensorFlow, and Git. Bonus: Hugging Face Transformers, Docker, "
        "AWS SageMaker. Fluent English required. Fully remote within Germany."
    ),
    "data_analyst_working_student_frankfurt.txt": (
        "Working Student Data Analyst (m/w/d), Frankfurt am Main. You will build "
        "dashboards and reports and derive insights from customer data. "
        "Requirements: enrolled in a quantitative field. Strong SQL, comfortable "
        "with Power BI or Tableau, and good Excel skills. Some Python for data "
        "cleaning is welcome. German at B2 level is required for this client-facing "
        "role; English is a plus. 20 hours per week, on-site."
    ),
}


# ------------------------------ Ingest -------------------------------------
def extract_file(path: str) -> str:
    """Read text from a PDF or plain-text file."""
    if path.lower().endswith(".pdf"):
        reader = PdfReader(path)
        return "\n".join((page.extract_text() or "") for page in reader.pages)
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        return f.read()


def chunk_text(text: str, source: str,
               size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP):
    """Split text into overlapping word windows, tagged with their source."""
    words = text.split()
    chunks, start = [], 0
    while start < len(words):
        end = start + size
        piece = " ".join(words[start:end]).strip()
        if piece:
            chunks.append({"text": piece, "source": source})
        if end >= len(words):
            break
        start += size - overlap
    return chunks


# ------------------------- Embed + index -----------------------------------
def build_index(chunks):
    """Embed chunks and build a cosine-similarity FAISS index."""
    texts = [c["text"] for c in chunks]
    emb = embedder.encode(texts, convert_to_numpy=True,
                          normalize_embeddings=True).astype(np.float32)
    index = faiss.IndexFlatIP(emb.shape[1])  # inner product on unit vectors = cosine
    index.add(emb)
    return index


# ------------------------------ Retrieve -----------------------------------
def retrieve(query, index, chunks, k=TOP_K):
    q = embedder.encode([query], convert_to_numpy=True,
                        normalize_embeddings=True).astype(np.float32)
    scores, idxs = index.search(q, min(k, len(chunks)))
    out = []
    for score, i in zip(scores[0], idxs[0]):
        if i != -1:
            out.append({**chunks[i], "score": float(score)})
    return out


# ------------------------------ Generate -----------------------------------
def generate_answer(query, retrieved, client):
    context = "\n\n".join(
        f"[{i+1}] (source: {r['source']})\n{r['text']}"
        for i, r in enumerate(retrieved)
    )
    system = (
        "You answer questions using ONLY the numbered context passages provided. "
        "Cite the passages you rely on inline with their bracket numbers, e.g. [1] "
        "or [2][3]. If the answer is not contained in the context, say you don't "
        "have enough information rather than guessing."
    )
    user = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer (with citations):"
    resp = client.chat.completions.create(
        model=GROQ_MODEL,
        messages=[{"role": "system", "content": system},
                  {"role": "user", "content": user}],
        temperature=0.2,
        max_tokens=600,
    )
    return resp.choices[0].message.content


# --------------------------- Gradio handlers -------------------------------
def do_build(files):
    """Build (or rebuild) the index from uploaded files, or the sample corpus."""
    all_chunks = []
    if files:
        for path in files:
            all_chunks += chunk_text(extract_file(path), os.path.basename(path))
        n_docs = len(files)
    else:
        for name, text in SAMPLE_DOCS.items():
            all_chunks += chunk_text(text, name)
        n_docs = len(SAMPLE_DOCS)

    if not all_chunks:
        return None, "No readable text found. Try a different file."

    index = build_index(all_chunks)
    state = {"index": index, "chunks": all_chunks}
    src = "uploaded document(s)" if files else "built-in sample job postings"
    return state, f"Indexed {len(all_chunks)} chunks from {n_docs} {src}. Ask away below."


def do_ask(query, state):
    if not state:
        return "Build the index first (the sample corpus loads automatically).", ""
    if not query.strip():
        return "Please enter a question.", ""

    retrieved = retrieve(query, state["index"], state["chunks"])
    sources_md = "\n\n".join(
        f"**[{i+1}]** *{r['source']}* — similarity {r['score']:.2f}\n\n> {r['text'][:300]}…"
        for i, r in enumerate(retrieved)
    )

    client = get_llm_client()
    if client is None:
        return (
            "**No `GROQ_API_KEY` set**, so generation is off — but retrieval works: "
            "the relevant passages are shown under *Retrieved sources*. Add a free "
            "Groq key as a Space secret to enable generated answers.",
            sources_md,
        )
    try:
        answer = generate_answer(query, retrieved, client)
    except Exception as e:  # surface API errors instead of crashing the UI
        answer = f"Generation error: {e}\n\n(The retrieved passages are still shown below.)"
    return answer, sources_md


# ------------------------------- UI ----------------------------------------
with gr.Blocks(title="RAG Document Q&A") as demo:
    gr.Markdown(
        "# 📄 RAG Document Q&A\n"
        "Ask questions and get answers **grounded in your documents**, with inline "
        "`[n]` citations. Upload your own PDFs/text, or just use the built-in "
        "sample job postings (loaded automatically).\n\n"
        "*Retrieval (embeddings + FAISS) runs locally; only the final answer is "
        "generated via Groq's free API.*"
    )
    state = gr.State()

    with gr.Row():
        files = gr.File(
            label="Upload PDF or .txt (optional — leave empty to use sample postings)",
            file_count="multiple", file_types=[".pdf", ".txt"], type="filepath",
        )
    build_btn = gr.Button("Build / Rebuild index", variant="secondary")
    status = gr.Markdown()

    query = gr.Textbox(label="Your question", lines=2,
                       placeholder="e.g. What skills do these data science roles have in common?")
    ask_btn = gr.Button("Ask", variant="primary")
    answer = gr.Markdown()
    with gr.Accordion("Retrieved sources", open=False):
        sources = gr.Markdown()

    gr.Examples(
        examples=[
            "What skills do these roles have in common?",
            "Which positions require cloud or AWS experience?",
            "Do any of these jobs need German language skills?",
            "Which role is the best fit for someone strong in PyTorch?",
        ],
        inputs=query,
    )

    build_btn.click(do_build, inputs=[files], outputs=[state, status])
    ask_btn.click(do_ask, inputs=[query, state], outputs=[answer, sources])
    # Auto-build the sample corpus on load so the demo is never an empty box.
    demo.load(lambda: do_build(None), outputs=[state, status])


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