ingest.py

"""
ingest.py
=========
Knowledge-base builder (roadmap Phase 2).

Pipeline:  raw PDFs -> extract text per page -> clean -> chunk
           -> embed -> store in ChromaDB with metadata (source, page, topic).

Run from the project root:

    python -m app.ingest                # ingest every PDF in data/raw_pdfs
    python -m app.ingest --reset        # wipe the collection first
    python -m app.ingest path/to.pdf    # ingest a single file
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Dict, List

import chromadb
from pypdf import PdfReader

from .config import settings
from .embeddings import get_embedder

# ---------------------------------------------------------------------------
# Topic tagging — lightweight keyword routing so retrieval can be filtered.
# ---------------------------------------------------------------------------
TOPIC_KEYWORDS: Dict[str, List[str]] = {
    "subnetting": ["subnet", "cidr", "netmask", "broadcast", "host range", "vlsm"],
    "tcp_ip": ["tcp", "udp", "ip header", "three-way", "handshake", "port", "segment"],
    "routing": ["router", "routing table", "ospf", "bgp", "static route", "next hop"],
    "switching_stp": ["switch", "vlan", "spanning tree", "stp", "mac address", "trunk"],
    "linux": ["ping", "traceroute", "ip addr", "netstat", "tcpdump", "ss ", "scp"],
    "iot": ["mqtt", "coap", "iot", "sensor", "broker", "publish", "subscribe", "edge"],
    "arp_l2": ["arp", "ethernet", "frame", "layer 2", "data link"],
}


def detect_topic(text: str) -> str:
    low = text.lower()
    best, best_hits = "general", 0
    for topic, kws in TOPIC_KEYWORDS.items():
        hits = sum(low.count(kw) for kw in kws)
        if hits > best_hits:
            best, best_hits = topic, hits
    return best


# ---------------------------------------------------------------------------
# Text cleaning + chunking
# ---------------------------------------------------------------------------
def clean_text(text: str) -> str:
    text = text.replace("\x00", " ")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def chunk_text(text: str, size: int, overlap: int) -> List[str]:
    """Word-based chunker (~tokens ≈ words). Overlap preserves context."""
    words = text.split()
    if not words:
        return []
    chunks: List[str] = []
    step = max(1, size - overlap)
    for start in range(0, len(words), step):
        chunk = " ".join(words[start : start + size]).strip()
        if chunk:
            chunks.append(chunk)
        if start + size >= len(words):
            break
    return chunks


def extract_pages(pdf_path: Path) -> List[str]:
    reader = PdfReader(str(pdf_path))
    return [clean_text(page.extract_text() or "") for page in reader.pages]


# ---------------------------------------------------------------------------
# ChromaDB
# ---------------------------------------------------------------------------
def get_collection():
    client = chromadb.PersistentClient(path=str(settings.VECTOR_DB_DIR))
    # We supply our own embeddings, so no embedding_function is configured.
    return client.get_or_create_collection(
        name=settings.CHROMA_COLLECTION, metadata={"hnsw:space": "cosine"}
    )


def reset_collection() -> None:
    client = chromadb.PersistentClient(path=str(settings.VECTOR_DB_DIR))
    try:
        client.delete_collection(settings.CHROMA_COLLECTION)
        print(f"[ingest] reset collection '{settings.CHROMA_COLLECTION}'")
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Main ingestion
# ---------------------------------------------------------------------------
def ingest_pdf(pdf_path: Path) -> int:
    embedder = get_embedder()
    collection = get_collection()

    ids: List[str] = []
    docs: List[str] = []
    metas: List[Dict] = []

    pages = extract_pages(pdf_path)
    processed_record = {"source": pdf_path.name, "chunks": []}

    for page_num, page_text in enumerate(pages, start=1):
        if not page_text:
            continue
        for ci, chunk in enumerate(chunk_text(page_text, settings.CHUNK_SIZE, settings.CHUNK_OVERLAP)):
            uid = f"{pdf_path.stem}_p{page_num}_c{ci}"
            topic = detect_topic(chunk)
            ids.append(uid)
            docs.append(chunk)
            metas.append(
                {"source": pdf_path.name, "page": page_num, "topic": topic, "chunk": ci}
            )
            processed_record["chunks"].append(
                {"id": uid, "page": page_num, "topic": topic, "text": chunk}
            )

    if not docs:
        print(f"[ingest] no extractable text in {pdf_path.name}")
        return 0

    vectors = embedder.encode(docs)
    collection.upsert(ids=ids, documents=docs, embeddings=vectors, metadatas=metas)

    # Persist cleaned chunks for inspection / debugging.
    out = settings.PROCESSED_DIR / f"{pdf_path.stem}.json"
    out.write_text(json.dumps(processed_record, indent=2))

    print(f"[ingest] {pdf_path.name}: {len(docs)} chunks  (embedder={embedder.name})")
    return len(docs)


def ingest_directory(directory: Path | None = None) -> int:
    directory = directory or settings.RAW_PDF_DIR
    total = 0
    pdfs = sorted(directory.glob("*.pdf"))
    if not pdfs:
        print(f"[ingest] no PDFs found in {directory}")
    for pdf in pdfs:
        total += ingest_pdf(pdf)
    print(f"[ingest] DONE — {total} chunks across {len(pdfs)} file(s).")
    return total


def main() -> None:
    parser = argparse.ArgumentParser(description="NetTutor AI ingestion")
    parser.add_argument("path", nargs="?", help="Single PDF to ingest (optional)")
    parser.add_argument("--reset", action="store_true", help="Reset collection first")
    args = parser.parse_args()

    if args.reset:
        reset_collection()

    if args.path:
        ingest_pdf(Path(args.path))
    else:
        ingest_directory()


if __name__ == "__main__":
    main()