#!/usr/bin/env python3
"""
ingest_corpus.py - Lab 4 setup
Ingests the synthetic corpus into a local ChromaDB vector store.
Each document chunk is tagged with source, trust level, and doc-id metadata.
Run inside the rag-harness container: python3 /lab/scripts/ingest_corpus.py
"""

import chromadb
import re
import sys

CORPUS_PATH = "/lab/seed/corpus.md"
CHROMA_PATH = "/lab/chroma_db"

def parse_corpus(path: str) -> list[dict]:
    """Parse the markdown corpus into document chunks with metadata."""
    with open(path) as f:
        content = f.read()

    sections = re.split(r"\n---\n", content)
    docs = []
    for section in sections:
        section = section.strip()
        if not section or section.startswith("#"):
            continue

        # Extract metadata lines
        doc_id_match = re.search(r"doc-id: ([\w-]+)", section)
        source_match = re.search(r"Source: (.+)", section)
        trust_match  = re.search(r"Trust: (\w+)", section)

        doc_id = doc_id_match.group(1)  if doc_id_match  else "unknown"
        source = source_match.group(1)  if source_match  else "unknown"
        trust  = trust_match.group(1)   if trust_match   else "unknown"

        # Strip metadata lines to get the body
        body_lines = [
            line for line in section.splitlines()
            if not line.startswith("Source:") and
               not line.startswith("Trust:") and
               not line.startswith("doc-id:") and
               not line.startswith("##")
        ]
        body = "\n".join(body_lines).strip()
        if not body:
            continue

        docs.append({
            "id":       doc_id,
            "text":     body,
            "metadata": {
                "source": source,
                "trust":  trust,
                "doc_id": doc_id,
            }
        })

    return docs


def main() -> None:
    docs = parse_corpus(CORPUS_PATH)
    print(f"Parsed {len(docs)} document chunks from corpus.")

    client     = chromadb.PersistentClient(path=CHROMA_PATH)
    collection = client.get_or_create_collection("lab19-rag")

    collection.add(
        documents = [d["text"]     for d in docs],
        ids       = [d["id"]       for d in docs],
        metadatas = [d["metadata"] for d in docs],
    )

    print(f"Ingested {len(docs)} chunks into ChromaDB at {CHROMA_PATH}")
    for d in docs:
        print(f"  [{d['metadata']['trust']:5s}] {d['id']}: {d['text'][:60].strip()}...")


if __name__ == "__main__":
    main()
