chunk_filings.py

"""
chunk_filings.py
----------------
Reads ParsedFiling JSON files produced by normalize_filing.py, splits each
section into sliding-window chunks, embeds them with OpenAI
text-embedding-3-small, and upserts into a local ChromaDB collection.

Usage:
    python chunk_filings.py output/parsed/SUPX_*.json
    python chunk_filings.py output/parsed/SUPX_*.json --db-path ./chroma_db --collection edgar_chunks
    python chunk_filings.py output/parsed/SUPX_*.json --dry-run   # counts chunks, no embedding

Environment variables:
    OPENAI_API_KEY   Required (for embedding calls).
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sqlite3
import sys
import time
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional

import tiktoken
import chromadb
import numpy as np
from openai import OpenAI

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("chunk_filings")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

EMBEDDING_MODEL   = "text-embedding-3-small"
EMBEDDING_DIM     = 1536          # fixed dim for text-embedding-3-small
MAX_TOKENS        = 512           # target chunk size (hard ceiling)
OVERLAP_PARAS     = 2             # paragraph overlap between consecutive chunks
BATCH_SIZE        = 64            # chunks per OpenAI embed call (API max = 2048)
COLLECTION_NAME   = "edgar_chunks"
EXISTING_ID_BATCH = 512           # per-batch lookup size when skipping existing chunks

# Use cl100k_base tokeniser (same as text-embedding-3-small)
_enc = None


def _get_encoder():
    global _enc
    if _enc is None:
        _enc = tiktoken.get_encoding("cl100k_base")
    return _enc


def _token_count(text: str) -> int:
    return len(_get_encoder().encode(text))


def _parse_filing_ts(filing_date: str) -> int:
    """
    Parse YYYY-MM-DD into Unix timestamp (UTC seconds). Returns 0 if invalid.
    """
    try:
        dt = datetime.strptime(filing_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
        return int(dt.timestamp())
    except Exception:
        return 0


def _infer_section_type(section_heading: str) -> str:
    s = (section_heading or "").strip().lower()
    if not s:
        return "other"
    if "cybersecurity" in s or "item 1c" in s:
        return "risk_cyber"
    if s in {"preamble", "signatures"}:
        return s
    if "securities and exchange commission" in s:
        return "preamble"
    if "table of contents" in s:
        return "toc"
    if "exhibit index" in s or "exhibits index" in s:
        return "exhibit_index"
    if s.startswith("item ") and "exhibit" in s:
        return "exhibit_index"
    if "risk factor" in s or "item 1a" in s:
        return "risk_factors"
    if "management discussion" in s or "md&a" in s or "results of operations" in s:
        return "mdna"
    if "financial statement" in s or "balance sheet" in s or "income statement" in s:
        return "financial_statements"
    if "beneficial ownership" in s or "schedule 13" in s:
        return "ownership"
    if "director" in s or "executive compensation" in s or "board" in s:
        return "governance"
    if "business" in s or "overview" in s:
        return "business"
    if "item " in s and ("8-k" in s or "6-k" in s):
        return "event"
    return "other"


def _is_boilerplate_text(section: str, text: str) -> bool:
    s = (section or "").strip().lower()
    t = (text or "").strip().lower()
    if s in {"preamble", "signatures"}:
        return True
    if "securities and exchange commission" in s:
        return True
    if "table of contents" in s:
        return True
    if "exhibit index" in s or "exhibits index" in s:
        return True
    if s.startswith("item ") and "exhibit" in s:
        return True
    if "securities and exchange commission" in t and "report of foreign private issuer" in t:
        return True
    if "pursuant to the requirements of the securities exchange act" in t:
        return True
    if "forward-looking statement" in s or "forward looking statement" in s:
        return True
    if "special note regarding" in s and ("forward-looking" in s or "forward looking" in s):
        return True
    if "special note regarding" in t and ("forward-looking" in t or "forward looking" in t):
        return True
    if "safe harbor" in t and ("forward-looking" in t or "forward looking" in t):
        return True
    return False


def _seed_topic_from_structure(form_type: str, section_type: str, section_heading: str) -> str:
    """
    High-precision structural labeling only (no fuzzy keyword scanning on chunk text).
    This keeps ingestion deterministic and lets embedding-based fallback handle ambiguous chunks.
    """
    f = (form_type or "").strip().upper()
    st = (section_type or "").strip().lower()
    sh = (section_heading or "").strip().lower()

    if st in {"preamble", "signatures", "toc", "exhibit_index"}:
        return "boilerplate"
    if st in {"risk_factors"}:
        return "risk_cyber"
    if st in {"risk_cyber"}:
        return "risk_cyber"
    if st in {"mdna", "financial_statements"}:
        return "financial"
    if st in {"governance"}:
        return "governance"
    if st in {"ownership"}:
        return "ownership"
    if st in {"business"}:
        return "business"
    if st in {"event"}:
        return "event"

    if f in {"8-K", "6-K", "SC TO"}:
        return "event"
    if "13G" in f or "13D" in f or f in {"3", "4", "5"}:
        return "ownership"
    if f == "DEF 14A":
        return "governance"
    if f in {"10-K", "10-Q", "20-F"}:
        if "item 1a" in sh:
            return "risk_cyber"
        if "item 1c" in sh or "cybersecurity" in sh:
            return "risk_cyber"
        if "item 7" in sh or "management discussion" in sh:
            return "financial"
        if "item 1" in sh or "item 4" in sh:
            return "business"

    return "general"


def _split_text_hard_by_tokens(text: str, max_tokens: int) -> List[str]:
    """
    Deterministic fallback splitter: slice by token count so every piece is <= max_tokens.
    """
    toks = _get_encoder().encode(text)
    if not toks:
        return []
    out: List[str] = []
    for i in range(0, len(toks), max_tokens):
        piece = _get_encoder().decode(toks[i : i + max_tokens]).strip()
        if piece:
            out.append(piece)
    return out


def _split_oversized_para(para: dict, max_tokens: int) -> List[dict]:
    """
    If a single paragraph exceeds max_tokens, split it into sub-paragraphs
    by double-newline first, then by sentence (period + space) if still too large.
    Returns a list of paragraph dicts (same schema, same metadata).
    """
    text = para["text"]
    if _token_count(text) <= max_tokens:
        return [para]

    # Try splitting on double-newline (most common in SEC filings)
    parts = [p.strip() for p in text.split("\n\n") if p.strip()]

    # If parts are still too large, split further on sentence boundaries
    fine_parts: List[str] = []
    for part in parts:
        if _token_count(part) > max_tokens:
            # Split on ". " keeping the period
            sentences = part.replace(". ", ".\n").split("\n")
            fine_parts.extend(s.strip() for s in sentences if s.strip())
        else:
            fine_parts.append(part)

    if not fine_parts:
        fine_parts = [text]

    # Final safety: hard-split any remaining oversized segment by tokens.
    split_parts: List[str] = []
    for part in fine_parts:
        if _token_count(part) <= max_tokens:
            split_parts.append(part)
        else:
            split_parts.extend(_split_text_hard_by_tokens(part, max_tokens))

    return [{**para, "text": part} for part in split_parts if _token_count(part) > 0]


def _deduplicate_paras(paragraphs: List[dict]) -> List[dict]:
    """
    Remove sec-parser oversized blob paragraphs that duplicate many neighbors.

    We intentionally avoid simple adjacent containment checks because normal
    prose can contain shorter adjacent phrases. Instead, we treat a paragraph
    as blob-like only when it is an outlier in length and substantially
    contains multiple other paragraphs.
    """
    if len(paragraphs) <= 1:
        return paragraphs

    texts = [str(p.get("text", "")).strip() for p in paragraphs]
    lengths = sorted(len(t) for t in texts if t)
    if not lengths:
        return paragraphs

    median_len = lengths[len(lengths) // 2]
    min_blob_len = max(250, int(median_len * 2.5))
    min_subpara_len = 40
    min_contained_paras = 2
    min_coverage_ratio = 0.20
    strong_outlier_ratio = 8.0

    drop_idx: set[int] = set()
    for i, text in enumerate(texts):
        if len(text) < min_blob_len:
            continue
        contained = 0
        covered_chars = 0
        for j, other in enumerate(texts):
            if i == j:
                continue
            if len(other) < min_subpara_len:
                continue
            if other == text:
                continue
            if other in text:
                contained += 1
                covered_chars += len(other)
        covered_enough = covered_chars >= int(len(text) * min_coverage_ratio)
        strong_outlier = len(text) >= int(max(1, median_len) * strong_outlier_ratio)
        if contained >= min_contained_paras and (covered_enough or strong_outlier):
            drop_idx.add(i)

    result: List[dict] = []
    seen_exact: set[str] = set()
    for i, para in enumerate(paragraphs):
        text = texts[i]
        if i in drop_idx:
            continue
        # Prevent exact duplicated paragraphs from entering the window builder.
        if text and text in seen_exact:
            continue
        if text:
            seen_exact.add(text)
        result.append(para)
    return result


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------

@dataclass
class Chunk:
    chunk_id:      str
    doc_id:        str
    ticker:        str
    form_type:     str
    filing_date:   str
    section:       str
    section_index: int
    chunk_index:   int
    char_offset:   int       # char_offset of the first paragraph in this chunk
    is_table:      bool      # True if ALL paragraphs in chunk are tables
    filing_ts:     int       # filing_date as unix timestamp for numeric date filters
    section_type:  str       # normalized section class used by retrieval model
    is_boilerplate: bool     # chunk-level boilerplate marker
    topic_tags:    str       # pipe-separated tags, e.g. "financial|risk"
    text:          str

    def metadata(self) -> dict:
        """Return flat dict suitable for ChromaDB metadata (strings/ints only)."""
        return {
            "doc_id":        self.doc_id,
            "ticker":        self.ticker,
            "form_type":     self.form_type,
            "filing_date":   self.filing_date,
            "section":       self.section,
            "section_index": self.section_index,
            "chunk_index":   self.chunk_index,
            "char_offset":   self.char_offset,
            "is_table":      int(self.is_table),   # Chroma wants int, not bool
            "filing_ts":     self.filing_ts,
            "section_type":  self.section_type,
            "is_boilerplate": int(self.is_boilerplate),
            "topic_tags":    self.topic_tags,
        }


# ---------------------------------------------------------------------------
# Chunking
# ---------------------------------------------------------------------------

def _make_chunk_id(doc_id: str, section_index: int, chunk_index: int) -> str:
    return f"{doc_id}__s{section_index:02d}_c{chunk_index:02d}"


def chunk_section(
    doc_id: str,
    ticker: str,
    form_type: str,
    filing_date: str,
    section_heading: str,
    section_index: int,
    paragraphs: list[dict],
    max_tokens: int = MAX_TOKENS,
    overlap: int = OVERLAP_PARAS,
) -> List[Chunk]:
    """
    Sliding-window chunker over a list of paragraphs.
    - Oversized single paragraphs are sentence-split before windowing.
    - Fills a window up to max_tokens, then emits a chunk.
    - Steps forward by (window_size - overlap) paragraphs.
    - Oversized elements are always split to satisfy max_tokens.
    """
    if not paragraphs:
        return []

    # 1. Deduplicate first (drop blob paras that contain all subsequent content)
    paragraphs = _deduplicate_paras(paragraphs)

    # 2. Expand any single paragraph that exceeds max_tokens into sub-paragraphs
    expanded: list[dict] = []
    for para in paragraphs:
        if _token_count(para["text"]) > max_tokens:
            expanded.extend(_split_oversized_para(para, max_tokens))
        else:
            expanded.append(para)
    paragraphs = expanded

    chunks: List[Chunk] = []
    chunk_idx = 0
    i = 0
    n = len(paragraphs)

    while i < n:
        window: list[dict] = []
        token_count = 0
        sep_tokens = _token_count("\n\n")

        j = i
        while j < n:
            para = paragraphs[j]
            para_tokens = _token_count(para["text"])
            add_sep = sep_tokens if window else 0

            # If adding this para would exceed limit AND we already have content,
            # close the current chunk first (except if it's the first para — must include it).
            if token_count + add_sep + para_tokens > max_tokens and window:
                break

            window.append(para)
            token_count += add_sep + para_tokens
            j += 1

        if not window:
            # Safety: a single paragraph still exceeds max_tokens, include alone.
            window = [paragraphs[i]]
            j = i + 1

        # Guardrail: enforce hard token cap in case separator interactions push over limit.
        text = "\n\n".join(p["text"] for p in window)
        while len(window) > 1 and _token_count(text) > max_tokens:
            window = window[:-1]
            text = "\n\n".join(p["text"] for p in window)

        if _token_count(text) > max_tokens:
            pieces = _split_text_hard_by_tokens(text, max_tokens)
            if not pieces:
                pieces = [text]
            first_para = window[0]
            section_type = _infer_section_type(section_heading)
            for piece in pieces:
                chunks.append(Chunk(
                    chunk_id      = _make_chunk_id(doc_id, section_index, chunk_idx),
                    doc_id        = doc_id,
                    ticker        = ticker,
                    form_type     = form_type,
                    filing_date   = filing_date,
                    section       = section_heading,
                    section_index = section_index,
                    chunk_index   = chunk_idx,
                    char_offset   = first_para.get("char_offset", 0),
                    is_table      = all(p.get("is_table", False) for p in window),
                    filing_ts     = _parse_filing_ts(filing_date),
                    section_type  = section_type,
                    is_boilerplate = _is_boilerplate_text(section_heading, piece),
                    topic_tags    = _seed_topic_from_structure(form_type, section_type, section_heading),
                    text          = piece,
                ))
                chunk_idx += 1
            step = max(1, len(window) - overlap)
            i += step
            continue

        all_tbl  = all(p.get("is_table", False) for p in window)
        first_para = window[0]
        section_type = _infer_section_type(section_heading)

        chunks.append(Chunk(
            chunk_id      = _make_chunk_id(doc_id, section_index, chunk_idx),
            doc_id        = doc_id,
            ticker        = ticker,
            form_type     = form_type,
            filing_date   = filing_date,
            section       = section_heading,
            section_index = section_index,
            chunk_index   = chunk_idx,
            char_offset   = first_para.get("char_offset", 0),
            is_table      = all_tbl,
            filing_ts     = _parse_filing_ts(filing_date),
            section_type  = section_type,
            is_boilerplate = _is_boilerplate_text(section_heading, text),
            topic_tags    = _seed_topic_from_structure(form_type, section_type, section_heading),
            text          = text,
        ))
        chunk_idx += 1

        # Slide forward: step = window_size - overlap, minimum 1
        step = max(1, len(window) - overlap)
        i += step

    return chunks


def chunks_from_filing(parsed: dict) -> List[Chunk]:
    """Extract all Chunks from a single ParsedFiling dict."""
    doc_id      = parsed["doc_id"]
    ticker      = parsed.get("ticker", "")
    form_type   = parsed.get("form_type", "")
    filing_date = parsed.get("filing_date", "")
    all_chunks: List[Chunk] = []

    for section in parsed.get("sections", []):
        paras = section.get("paragraphs", [])
        # Skip trivially empty sections
        if not paras:
            continue
        section_chunks = chunk_section(
            doc_id          = doc_id,
            ticker          = ticker,
            form_type       = form_type,
            filing_date     = filing_date,
            section_heading = section["heading"],
            section_index   = section["section_index"],
            paragraphs      = paras,
        )
        all_chunks.extend(section_chunks)

    return all_chunks


# ---------------------------------------------------------------------------
# Ingestion metadata policy
# ---------------------------------------------------------------------------
# Open-core intentionally keeps deterministic structural labels only:
# - section_type
# - is_boilerplate
# - topic_tags seeded from structure

# ---------------------------------------------------------------------------
# Embedding (OpenAI only)
# ---------------------------------------------------------------------------


class Embedder:
    """
    OpenAI embedding client wrapper.

    Usage:
        e = Embedder()                 # requires OPENAI_API_KEY with credits
        vecs = e.embed(["text1", "text2"])
    """

    def __init__(self, openai_api_key: Optional[str] = None) -> None:
        api_key = openai_api_key or os.environ.get("OPENAI_API_KEY", "")
        if not api_key:
            raise ValueError("OPENAI_API_KEY is not set.")
        self._oai = OpenAI(api_key=api_key)
        log.info(f"Embedder: OpenAI {EMBEDDING_MODEL}")

    def embed(self, texts: List[str]) -> List[List[float]]:
        return self._embed_openai(texts)

    def _embed_openai(self, texts: List[str]) -> List[List[float]]:
        while True:
            try:
                resp = self._oai.embeddings.create(input=texts, model=EMBEDDING_MODEL)
                return [item.embedding for item in resp.data]
            except Exception as e:
                if "rate" in str(e).lower():
                    log.warning(f"Rate-limited, waiting 10s… ({e})")
                    time.sleep(10)
                else:
                    raise

def embed_all(chunks: List[Chunk], embedder: Embedder) -> List[List[float]]:
    """Embed all chunks in batches."""
    embeddings: List[List[float]] = []
    total = len(chunks)

    for start in range(0, total, BATCH_SIZE):
        batch = chunks[start : start + BATCH_SIZE]
        texts = [c.text for c in batch]
        log.info(f"  Embedding batch {start}–{min(start + BATCH_SIZE, total) - 1} / {total - 1}")
        embeddings.extend(embedder.embed(texts))

    return embeddings


# ---------------------------------------------------------------------------
# ChromaDB upsert
# ---------------------------------------------------------------------------

def upsert_to_chroma(
    chunks: List[Chunk],
    embeddings: List[List[float]],
    collection: chromadb.Collection,
) -> None:
    """Upsert chunks + embeddings into ChromaDB (idempotent by chunk_id)."""
    if not chunks:
        return

    ids       = [c.chunk_id   for c in chunks]
    documents = [c.text        for c in chunks]
    metadatas = [c.metadata()  for c in chunks]

    collection.upsert(
        ids        = ids,
        embeddings = embeddings,
        documents  = documents,
        metadatas  = metadatas,
    )
    log.info(f"  Upserted {len(chunks)} chunks into collection.")


def init_sparse_index_db(db_path: str) -> None:
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS chunk_store (
                chunk_id TEXT PRIMARY KEY,
                doc_id TEXT,
                ticker TEXT,
                form_type TEXT,
                filing_date TEXT,
                filing_ts INTEGER,
                section TEXT,
                section_index INTEGER,
                chunk_index INTEGER,
                char_offset INTEGER,
                is_table INTEGER,
                section_type TEXT,
                is_boilerplate INTEGER,
                topic_tags TEXT,
                text TEXT
            )
            """
        )
        conn.execute("CREATE INDEX IF NOT EXISTS idx_chunk_store_ticker ON chunk_store(ticker)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_chunk_store_form ON chunk_store(form_type)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_chunk_store_date ON chunk_store(filing_date)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_chunk_store_ts ON chunk_store(filing_ts)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_chunk_store_doc ON chunk_store(doc_id)")
        conn.execute(
            """
            CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts
            USING fts5(chunk_id UNINDEXED, text, tokenize='unicode61')
            """
        )
        conn.commit()
    finally:
        conn.close()


def upsert_to_sparse_index(chunks: List[Chunk], db_path: str) -> None:
    if not chunks:
        return

    conn = sqlite3.connect(db_path)
    try:
        conn.execute("BEGIN")
        conn.executemany(
            """
            INSERT OR REPLACE INTO chunk_store (
                chunk_id, doc_id, ticker, form_type, filing_date, filing_ts,
                section, section_index, chunk_index, char_offset, is_table,
                section_type, is_boilerplate, topic_tags, text
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            [
                (
                    c.chunk_id,
                    c.doc_id,
                    c.ticker,
                    c.form_type,
                    c.filing_date,
                    c.filing_ts,
                    c.section,
                    c.section_index,
                    c.chunk_index,
                    c.char_offset,
                    int(c.is_table),
                    c.section_type,
                    int(c.is_boilerplate),
                    c.topic_tags,
                    c.text,
                )
                for c in chunks
            ],
        )
        # FTS5 rows are removed by rowid to avoid stale duplicates when a chunk
        # is re-upserted with changed text.
        conn.executemany(
            "DELETE FROM chunk_fts WHERE rowid IN (SELECT rowid FROM chunk_fts WHERE chunk_id = ?)",
            [(c.chunk_id,) for c in chunks],
        )
        conn.executemany(
            "INSERT INTO chunk_fts (chunk_id, text) VALUES (?, ?)",
            [(c.chunk_id, c.text) for c in chunks],
        )
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def upsert_sparse_for_run(
    *,
    chunks: List[Chunk],
    index_chunks: List[Chunk],
    db_path: str,
    force: bool,
) -> int:
    """
    Upsert sparse index payload for one ingestion iteration.
    - force=True  -> refresh all chunks
    - force=False -> only refresh newly indexed chunks
    """
    payload = chunks if force else index_chunks
    if not payload:
        return 0
    upsert_to_sparse_index(payload, db_path)
    return len(payload)


def existing_ids_in_collection(
    collection: chromadb.Collection,
    ids: List[str],
    batch_size: int = EXISTING_ID_BATCH,
) -> set[str]:
    """
    Return the subset of ids already present in ChromaDB.
    """
    if not ids:
        return set()

    found: set[str] = set()
    for i in range(0, len(ids), batch_size):
        batch = ids[i : i + batch_size]
        try:
            res = collection.get(ids=batch, include=["metadatas"])
        except Exception:
            continue
        batch_ids = res.get("ids") or []
        found.update(str(x) for x in batch_ids if x)
    return found


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    ap = argparse.ArgumentParser(description="Chunk + embed EDGAR ParsedFiling JSONs into ChromaDB.")
    ap.add_argument("files", nargs="+", help="ParsedFiling JSON files")
    DB_DIR = str(Path(__file__).parent / "data" / "chroma_db")
    SPARSE_DB = str(Path(__file__).parent / "data" / "sparse_index.db")
    ap.add_argument("--db-path",    default=DB_DIR,   help="ChromaDB persist directory")
    ap.add_argument("--sparse-db",  default=SPARSE_DB, help="SQLite FTS5 sparse index database path")
    ap.add_argument("--collection", default=COLLECTION_NAME, help="ChromaDB collection name")
    ap.add_argument("--model",      default="openai", choices=["openai"], help="Embedding backend (OpenAI only)")
    ap.add_argument("--dry-run",    action="store_true",     help="Count chunks without embedding")
    ap.add_argument("--force",      action="store_true",
                    help="Re-embed and upsert all chunks, including existing chunk_ids.")
    args = ap.parse_args()

    # ChromaDB
    db = chromadb.PersistentClient(path=args.db_path)
    collection = db.get_or_create_collection(
        name=args.collection,
        metadata={"hnsw:space": "cosine"},
    )
    log.info(f"ChromaDB: {args.db_path}  collection={args.collection}  existing={collection.count()}")
    if not args.dry_run:
        init_sparse_index_db(args.sparse_db)
        log.info(f"Sparse index: {args.sparse_db}")

    # Embedder
    embedder: Optional[Embedder] = None
    if not args.dry_run:
        embedder = Embedder()

    total_chunks = 0
    total_upserted = 0
    total_skipped_existing = 0

    for fp in args.files:
        path = Path(fp)
        log.info(f"Processing: {path.name}")

        with open(path, encoding="utf-8") as fh:
            parsed = json.load(fh)

        chunks = chunks_from_filing(parsed)
        log.info(f"  → {len(chunks)} chunks  sections={len(parsed.get('sections', []))}")
        total_chunks += len(chunks)

        if args.dry_run:
            for c in chunks[:3]:
                log.info(f"    [dry] {c.chunk_id}  {_token_count(c.text)} tok  '{c.section[:50]}'")
            continue

        if not chunks:
            continue

        index_chunks = chunks
        if not args.force:
            ids = [c.chunk_id for c in chunks]
            existing_ids = existing_ids_in_collection(collection, ids)
            if existing_ids:
                index_chunks = [c for c in chunks if c.chunk_id not in existing_ids]
                skipped_existing = len(chunks) - len(index_chunks)
                total_skipped_existing += skipped_existing
                log.info(f"  Skipped {skipped_existing} existing chunks (already indexed).")

        if not index_chunks:
            continue

        embeddings = embed_all(index_chunks, embedder)
        upsert_to_chroma(index_chunks, embeddings, collection)
        # Keep sparse index in sync without downgrading skipped chunk metadata.
        upsert_sparse_for_run(
            chunks=chunks,
            index_chunks=index_chunks,
            db_path=args.sparse_db,
            force=args.force,
        )
        total_upserted += len(index_chunks)

    log.info(f"Done. Total chunks: {total_chunks}. "
             f"Newly upserted: {total_upserted}. "
             f"Skipped existing: {total_skipped_existing}. "
             f"Collection size: {collection.count() if not args.dry_run else 'N/A (dry-run)'}.")


if __name__ == "__main__":
    main()