#!/usr/bin/env python3
"""
rag_weak.py - Lab 4, Step 2: the VULNERABLE RAG
Retrieves from ChromaDB (no trust filtering) and passes all chunks
including attacker-injected ones into the Ollama model context.
Demonstrates indirect prompt injection and cross-tenant/access-bypass.

Usage (inside rag-harness container):
    python3 /lab/scripts/rag_weak.py "How do I reset my password?"
"""

import sys
import chromadb
import ollama
import os

CHROMA_PATH  = "/lab/chroma_db"
OLLAMA_HOST  = os.environ.get("OLLAMA_HOST", "http://host.docker.internal:11434")
MODEL        = os.environ.get("RAG_MODEL", "llama3.2")
TOP_K        = 3


def retrieve(query: str, collection) -> list[dict]:
    results = collection.query(query_texts=[query], n_results=TOP_K)
    chunks  = []
    for i, doc in enumerate(results["documents"][0]):
        meta = results["metadatas"][0][i]
        chunks.append({"text": doc, "meta": meta})
    return chunks


def answer_weak(query: str, chunks: list[dict]) -> str:
    """No trust filtering - ALL retrieved chunks go into context."""
    context = "\n\n---\n\n".join(c["text"] for c in chunks)
    prompt  = (
        "You are a helpful IT support assistant. "
        "Use the following retrieved documents to answer the user's question.\n\n"
        f"DOCUMENTS:\n{context}\n\n"
        f"USER QUESTION: {query}\n\nANSWER:"
    )

    client   = ollama.Client(host=OLLAMA_HOST)
    response = client.generate(model=MODEL, prompt=prompt)
    return response["response"]


def main() -> None:
    query = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else "How do I reset my password?"

    client     = chromadb.PersistentClient(path=CHROMA_PATH)
    collection = client.get_collection("lab19-rag")
    chunks     = retrieve(query, collection)

    print(f"=== WEAK RAG: query={query!r} ===")
    print(f"Retrieved {len(chunks)} chunks:")
    for c in chunks:
        trust = c["meta"].get("trust", "?")
        doc_id = c["meta"].get("doc_id", "?")
        print(f"  [{trust:5s}] {doc_id}: {c['text'][:80].strip()}...")
    print()
    answer = answer_weak(query, chunks)
    print("=== MODEL RESPONSE ===")
    print(answer)


if __name__ == "__main__":
    main()
