"""Build the FAISS index from a folder of .txt medical reference docs.
Run once locally: python -m app.build_index --docs ./docs --out ./rag_index
"""
import argparse, json, os
from pathlib import Path
import faiss, numpy as np
from sentence_transformers import SentenceTransformer

def chunk_text(text, size=500, overlap=50):
    words = text.split()
    for i in range(0, len(words), size - overlap):
        yield " ".join(words[i:i+size])

def main(docs_dir, out_dir):
    encoder = SentenceTransformer("pritamdeka/S-PubMedBert-MS-MARCO")
    chunks = []
    for path in Path(docs_dir).glob("*.txt"):
        text = path.read_text()
        for chunk in chunk_text(text):
            chunks.append({"source": path.name, "text": chunk})

    embeddings = encoder.encode(
        [c["text"] for c in chunks],
        normalize_embeddings=True,
        show_progress_bar=True,
    )
    index = faiss.IndexFlatIP(embeddings.shape[1])
    index.add(np.asarray(embeddings, dtype="float32"))

    Path(out_dir).mkdir(exist_ok=True)
    faiss.write_index(index, f"{out_dir}/faiss.index")
    with open(f"{out_dir}/chunks.jsonl", "w") as f:
        for c in chunks:
            f.write(json.dumps(c) + "\n")
    print(f"Indexed {len(chunks)} chunks → {out_dir}")

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--docs", dest="docs_dir", required=True)
    p.add_argument("--out", dest="out_dir", default="./rag_index")
    main(**vars(p.parse_args()))