"""
index_findings.py — STEP 2: Build the vector database from your findings.

Run this ONCE (and again only when you add many new findings).
It reads corpus.txt, splits it into individual findings, turns each into a
vector with a local embedding model, and stores them in ChromaDB on disk.

Run:
    source venv/bin/activate
    python index_findings.py
"""

import re
import chromadb
from sentence_transformers import SentenceTransformer

CORPUS_PATH = "corpus.txt"
DB_PATH = "./nc_vectors"          # folder where the vector DB is stored
COLLECTION = "findings"

# ---- 1. Read the corpus ----
with open(CORPUS_PATH, "r", encoding="utf-8") as f:
    raw = f.read()
print(f"Read {len(raw)} characters from {CORPUS_PATH}")

# ---- 2. Split into individual findings ----
# Your corpus uses <|nc|> ... <|end|> markers around each finding.
# We split on <|nc|> and keep each block as one finding.
blocks = [b.strip() for b in raw.split("<|nc|>") if b.strip()]
# Clean the trailing <|end|> marker from each block.
findings = []
for b in blocks:
    text = b.replace("<|end|>", "").strip()
    if len(text) > 20:            # skip tiny/empty fragments
        findings.append(text)
print(f"Found {len(findings)} findings to index")

if not findings:
    raise SystemExit("No findings found — check corpus.txt format / markers.")

# ---- 3. Load the embedding model (downloads once, then cached locally) ----
print("Loading embedding model (first run downloads ~90MB)...")
embedder = SentenceTransformer("all-MiniLM-L6-v2")

# ---- 4. Create / reset the ChromaDB collection ----
client = chromadb.PersistentClient(path=DB_PATH)
# Start fresh each time you re-index
try:
    client.delete_collection(COLLECTION)
except Exception:
    pass
col = client.create_collection(COLLECTION)

# ---- 5. Embed and store, in batches (faster) ----
BATCH = 64
for start in range(0, len(findings), BATCH):
    chunk = findings[start:start + BATCH]
    vectors = embedder.encode(chunk).tolist()
    ids = [str(i) for i in range(start, start + len(chunk))]
    col.add(ids=ids, embeddings=vectors, documents=chunk)
    print(f"  indexed {start + len(chunk)} / {len(findings)}")

print(f"\nDone. Vector database saved to {DB_PATH}")
print(f"Total findings indexed: {col.count()}")