cli.py

"""
Vector CLI — JSON-first interface designed for AI agent consumption.

Commands
--------
ingest      Convert a PDF, store it, chunk it, and embed it.
query       Semantic search across one or more documents.
ask         Query expansion + search + AI-generated answer.
list        List ingested documents (with optional name/tag filter).
status      Show collection health (doc count, chunk count, issues).
repair      Find and optionally fix incomplete ingestions.
rename      Rename a document in the manifest and vector store.
tag         Set, update, or remove tags on a document.
tags        List all tag keys and their distinct values across the collection.
collections List all collections in the data directory.
delete      Remove a document from the document store and vector store.
figures     List or export extracted figures for a document.

All output is newline-terminated JSON printed to stdout.
Long-running commands (ingest, repair --fix) emit intermediate progress
lines before the final result line.
Errors are printed to stderr as {"error": "...", "error_code": "..."} with exit code 1.

Error codes
-----------
    conversion_failed   PDF could not be parsed
    store_failed        Could not write to doc store
    chunking_failed     Chunking step failed
    embedding_failed    ChromaDB write failed
    query_failed        Vector search failed
    not_found           --name filter matched no documents
    list_failed         Could not read store
    delete_failed       Could not delete document
    status_failed       Could not read store for status check

Data layout
-----------
    <data_dir>/<collection>/manifest.json
    <data_dir>/<collection>/documents/<hash>/source.pdf
    <data_dir>/<collection>/documents/<hash>/figures/pic_000.png, tbl_000.png
    <data_dir>/<collection>/embeddings/  (ChromaDB persist directory)

Usage examples
--------------
    python cli.py ingest data/report.pdf
    python cli.py ingest data/report.pdf --extract-figures --tag type=report --tag jurisdiction=CA
    python cli.py query "What are the main AI models?" --top-k 3 --window 1
    python cli.py query "tables" --name "coweta" --tag type=report
    python cli.py query "tables" --file-hash 11465328351749295394
    python cli.py list
    python cli.py list --name report
    python cli.py list --tag type=report
    python cli.py tag --file-hash 11465328351749295394 author=IBM type=whitepaper
    python cli.py tag --name "report" author=IBM
    python cli.py tag --file-hash 11465328351749295394 --remove state
    python cli.py tags
    python cli.py query "setbacks" --tag state=GA,FL
    python cli.py status
    python cli.py repair
    python cli.py repair --fix
    python cli.py delete 11465328351749295394
"""

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, Optional

from dotenv import load_dotenv
load_dotenv(override=False)

from agents import expand_query, answer_question
from chunker import Chunker
from converter import Converter
from document_store import DocumentStore
from models import FigureRecord
from rendering import crop_figure
from vector_store import VectorStore
import r2 as r2_mod

DEFAULT_DATA_DIR = os.environ.get("VECTOR_DATA_DIR", "./vector_data")
DEFAULT_COLLECTION = "vector"


# ---------------------------------------------------------------------------
# Output helpers
# ---------------------------------------------------------------------------

def _ok(payload: Any) -> None:
    print(json.dumps(payload, default=str), flush=True)


def _progress(payload: Any) -> None:
    """Emit a progress line to stdout during long-running operations.

    All progress lines include ``"type": "progress"`` so agents can distinguish
    them from final result lines (which have no ``type`` field).
    """
    if isinstance(payload, dict):
        payload = {"type": "progress", **payload}
    print(json.dumps(payload, default=str), flush=True)


def _err(message: str, error_code: str = "error") -> None:
    print(json.dumps({"error": message, "error_code": error_code}), file=sys.stderr)
    sys.exit(1)


def _normalize_tag_value(v: str) -> str:
    """Normalize a tag value: lowercase, strip whitespace, replace spaces with hyphens."""
    return v.strip().lower().replace(" ", "-")


def _normalize_tag_key(k: str) -> str:
    """Normalize a tag key: lowercase, strip whitespace."""
    return k.strip().lower()


def _parse_tags(raw: list, normalize: bool = False) -> Dict[str, str]:
    """Parse a list of ``key=value`` strings into a dict.

    When *normalize* is True, keys and values are lowercased and spaces in
    values are replaced with hyphens.  A warning is printed to stderr if a
    value required normalization (contained uppercase or spaces).

    Values that contain commas are returned as lists to express OR semantics
    in ChromaDB ``$in`` queries.  Normalization is applied to each element
    of a comma-separated value.

    Raises:
        SystemExit: If any entry cannot be split on ``=``.
    """
    if not raw:
        return {}
    tags: Dict = {}
    for item in raw:
        if "=" not in item:
            _err(f"Invalid tag format (expected key=value): {item!r}", "invalid_tag")
        k, _, v = item.partition("=")
        k = k.strip()
        v = v.strip()

        # OR logic: comma-separated values become a list
        if "," in v:
            parts = [p.strip() for p in v.split(",") if p.strip()]
            if normalize:
                normalized_parts = [_normalize_tag_value(p) for p in parts]
                for orig, norm in zip(parts, normalized_parts):
                    if orig != norm:
                        print(
                            f"Warning: tag value {orig!r} normalized to {norm!r}",
                            file=sys.stderr,
                        )
                parts = normalized_parts
            norm_k = _normalize_tag_key(k)
            tags[norm_k] = parts
        else:
            norm_k = _normalize_tag_key(k)
            norm_v = _normalize_tag_value(v) if normalize else v
            if normalize and (v != norm_v or k != norm_k):
                if v != norm_v:
                    print(
                        f"Warning: tag value {v!r} normalized to {norm_v!r}",
                        file=sys.stderr,
                    )
            tags[norm_k] = norm_v
    return tags


def _stores(args: argparse.Namespace):
    """Return (DocumentStore, VectorStore) rooted under <data_dir>/<collection>/.

    Directory layout:
        <data_dir>/<collection>/manifest.json          (DocumentStore manifest)
        <data_dir>/<collection>/documents/<hash>/      (per-document assets)
        <data_dir>/<collection>/embeddings/            (ChromaDB persist dir)
    """
    root = Path(args.data_dir) / args.collection
    doc_store = DocumentStore(str(root))
    vs = VectorStore(
        collection_name=args.collection,
        persist_directory=str(root / "embeddings"),
    )
    return doc_store, vs


def _extract_and_save_figures(chunks, doc_store, file_hash):
    """Extract figure images from the stored PDF using figure refs on chunks.

    Crops each referenced figure from the PDF using its bounding box and saves
    it as a PNG. Returns the count of figures extracted.
    """
    # Collect unique figure refs across all chunks
    seen_refs = set()
    all_figures = []
    chunk_map = {}  # ref -> list of chunk indices

    for chunk in chunks:
        for fig_ref in chunk.figures:
            if fig_ref.ref not in chunk_map:
                chunk_map[fig_ref.ref] = []
            chunk_map[fig_ref.ref].append(chunk.index)
            if fig_ref.ref not in seen_refs:
                seen_refs.add(fig_ref.ref)
                all_figures.append(fig_ref)

    if not all_figures:
        return 0

    try:
        pdf_path = str(doc_store.get_pdf_path(file_hash))
    except FileNotFoundError:
        return 0

    figure_records = []
    images = {}
    for i, fig_ref in enumerate(all_figures):
        kind_prefix = "pic" if fig_ref.kind == "picture" else "tbl"
        fig_id = f"{file_hash}_{kind_prefix}_{i:03d}"

        record = FigureRecord(
            id=fig_id,
            file_hash=file_hash,
            ref=fig_ref.ref,
            kind=fig_ref.kind,
            page_number=fig_ref.page_number,
            bbox=fig_ref.bbox,
            caption=fig_ref.caption,
            nearby_chunk_indices=chunk_map.get(fig_ref.ref, []),
        )
        figure_records.append(record)

        # Crop the figure from the PDF if we have a bbox
        if fig_ref.bbox is not None:
            try:
                img = crop_figure(pdf_path, fig_ref.bbox)
                images[fig_id] = img
            except Exception:
                pass  # Skip figures that fail to render

    doc_store.save_figures(file_hash, figure_records, images)
    return len(figure_records)


# ---------------------------------------------------------------------------
# R2 helper
# ---------------------------------------------------------------------------

def _r2_upload_and_tag(args, doc_store, vs, file_hash: str, has_figures: bool) -> dict:
    """Upload PDF (and optionally figures) to R2, then write tags to manifest + vector store.

    Returns the dict of new tags that were written.
    Calls _err() (sys.exit) on a fatal upload failure.
    Figure upload failures are non-fatal: emitted as warning progress lines.
    """
    if not r2_mod.is_configured():
        _err(
            "R2 not configured — set R2_ACCOUNT_ID, R2_BUCKET_NAME, "
            "R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY",
            "r2_not_configured",
        )
    _progress({"status": "uploading_r2", "file_hash": file_hash})
    new_tags: dict = {}

    try:
        pdf_path = doc_store.get_pdf_path(file_hash)
        key = r2_mod.upload_pdf(file_hash, pdf_path)
        new_tags["r2_pdf_key"] = key
        pub = r2_mod.public_url(key)
        if pub:
            new_tags["r2_pdf_url"] = pub
    except Exception as e:
        _err(f"R2 PDF upload failed: {e}", "r2_upload_failed")

    if has_figures:
        try:
            figures_dir = doc_store._figures_dir(file_hash)
            if figures_dir.exists():
                keys = r2_mod.upload_figures(file_hash, figures_dir)
                if keys:
                    new_tags["r2_figures_key"] = f"figures/{file_hash}/"
                    pub_base = r2_mod.public_url(f"figures/{file_hash}/")
                    if pub_base:
                        new_tags["r2_figures_url"] = pub_base
        except Exception as e:
            print(
                json.dumps({"type": "warning", "message": f"R2 figures upload failed: {e}"}),
                flush=True,
            )

    try:
        doc_store.update_tags(file_hash, new_tags)
        vs.set_tag_metadata(file_hash, new_tags)
    except Exception as e:
        print(
            json.dumps({"type": "warning", "message": f"R2 tag write failed: {e}"}),
            flush=True,
        )

    return new_tags


# ---------------------------------------------------------------------------
# Command handlers
# ---------------------------------------------------------------------------

def cmd_ingest(args: argparse.Namespace) -> None:
    if not Path(args.file).exists():
        _err(f"File not found: {args.file}", "not_found")

    tags = _parse_tags(getattr(args, "tag", None) or [])

    converter = Converter()
    doc_store, vs = _stores(args)

    if args.chunk_size > 0:
        _cmd_ingest_chunked(args, converter, doc_store, vs, tags=tags)
        return

    _progress({"status": "converting", "file": args.file})
    try:
        dl_doc = converter.convert(args.file)
    except Exception as e:
        _err(f"Conversion failed: {e}", "conversion_failed")

    file_hash = str(dl_doc.origin.binary_hash)
    if doc_store.exists(file_hash):
        if not args.force:
            _ok({
                "status": "already_ingested",
                "collection": args.collection,
                "file_hash": file_hash,
                "document_name": dl_doc.name,
                "filename": dl_doc.origin.filename,
            })
            return
        # --force: delete existing data and re-ingest
        doc_store.delete(file_hash)
        vs.delete(file_hash)

    _progress({"status": "storing", "file_hash": file_hash, "page_count": len(dl_doc.pages)})
    try:
        file_hash = doc_store.create(dl_doc, source_pdf_path=args.file)
    except Exception as e:
        _err(f"Document store failed: {e}", "store_failed")

    # Save the DoclingDocument as JSON for future use (non-chunked mode only;
    # for chunked mode the PDF is preserved for re-conversion).
    try:
        docling_path = doc_store.get_docling_doc_path(file_hash)
        docling_path.parent.mkdir(parents=True, exist_ok=True)
        with open(str(docling_path), "w", encoding="utf-8") as f:
            f.write(dl_doc.model_dump_json(indent=2))
    except Exception:
        pass  # Non-fatal — PDF is always available for re-conversion

    _progress({"status": "chunking"})
    try:
        chunker = Chunker(dl_doc)
        chunks = chunker.chunk()
    except Exception as e:
        doc_store.set_status(file_hash, "failed")
        _err(f"Chunking failed: {e}", "chunking_failed")

    _progress({"status": "embedding", "chunk_count": len(chunks)})
    try:
        vs.create(chunks, tags=tags if tags else None)
        doc_store.set_status(file_hash, "complete")
        doc_store.update_chunk_count(file_hash, len(chunks))
    except Exception as e:
        doc_store.set_status(file_hash, "failed")
        _err(f"Embedding failed: {e}", "embedding_failed")

    # Store tags in manifest
    if tags:
        try:
            doc_store.update_tags(file_hash, tags)
        except Exception:
            pass  # Non-fatal

    # Extract figures if requested
    figure_count = 0
    if args.extract_figures:
        _progress({"status": "extracting_figures"})
        figure_count = _extract_and_save_figures(chunks, doc_store, file_hash)

    r2_tags: dict = {}
    if r2_mod.is_configured():
        r2_tags = _r2_upload_and_tag(args, doc_store, vs, file_hash, figure_count > 0)

    _ok({
        "status": "ingested",
        "collection": args.collection,
        "file_hash": file_hash,
        "document_name": dl_doc.name,
        "filename": dl_doc.origin.filename,
        "page_count": len(dl_doc.pages),
        "chunk_count": len(chunks),
        "figure_count": figure_count,
        "tags": tags,
        **r2_tags,
    })


def _cmd_ingest_chunked(
    args: argparse.Namespace,
    converter: Converter,
    doc_store: DocumentStore,
    vs: VectorStore,
    tags: Optional[Dict[str, str]] = None,
) -> None:
    """Ingest a PDF by converting it in sequential page-range chunks.

    Each chunk is converted, chunked, and embedded independently so that peak
    memory never exceeds what a single page range requires.  Chunk indices are
    renumbered globally so that window-based context retrieval works correctly
    across range boundaries.
    """
    total_pages = converter.page_count(args.file)
    chunk_size = args.chunk_size
    ranges = [
        (start, min(start + chunk_size - 1, total_pages))
        for start in range(1, total_pages + 1, chunk_size)
    ]
    total_ranges = len(ranges)

    # Convert the first range to obtain document identity before storing.
    start, end = ranges[0]
    _progress({
        "status": "converting",
        "file": args.file,
        "page_range": f"{start}-{end}",
        "range_index": 1,
        "total_ranges": total_ranges,
        "total_pages": total_pages,
    })
    try:
        first_dl_doc = converter.convert_page_range(args.file, start, end)
    except Exception as e:
        _err(f"Conversion failed on pages {start}-{end}: {e}", "conversion_failed")

    file_hash = str(first_dl_doc.origin.binary_hash)
    if doc_store.exists(file_hash):
        if not args.force:
            _ok({
                "status": "already_ingested",
                "collection": args.collection,
                "file_hash": file_hash,
                "document_name": first_dl_doc.name,
                "filename": first_dl_doc.origin.filename,
            })
            return
        # --force: delete existing data and re-ingest
        doc_store.delete(file_hash)
        vs.delete(file_hash)

    _progress({"status": "storing", "file_hash": file_hash, "page_count": total_pages})
    try:
        file_hash = doc_store.create(
            first_dl_doc, source_pdf_path=args.file, page_count=total_pages
        )
    except Exception as e:
        _err(f"Document store failed: {e}", "store_failed")

    chunk_offset = 0
    total_chunks = 0
    all_chunks_for_figures = []  # Collect for figure extraction
    for i, (start, end) in enumerate(ranges):
        dl_doc = first_dl_doc if i == 0 else None
        if dl_doc is None:
            _progress({
                "status": "converting",
                "file": args.file,
                "page_range": f"{start}-{end}",
                "range_index": i + 1,
                "total_ranges": total_ranges,
            })
            try:
                dl_doc = converter.convert_page_range(args.file, start, end)
            except Exception as e:
                doc_store.set_status(file_hash, "failed")
                _err(f"Conversion failed on pages {start}-{end}: {e}", "conversion_failed")

        _progress({
            "status": "chunking",
            "page_range": f"{start}-{end}",
            "range_index": i + 1,
            "total_ranges": total_ranges,
        })
        try:
            range_chunks = Chunker(dl_doc).chunk()
        except Exception as e:
            doc_store.set_status(file_hash, "failed")
            _err(f"Chunking failed on pages {start}-{end}: {e}", "chunking_failed")

        # Renumber indices to be globally sequential across all ranges so that
        # window-based context retrieval works correctly for cross-range queries.
        for chunk in range_chunks:
            new_index = chunk_offset + chunk.index
            chunk.index = new_index
            chunk.id = f"{file_hash}_{new_index}"

        _progress({
            "status": "embedding",
            "page_range": f"{start}-{end}",
            "range_index": i + 1,
            "total_ranges": total_ranges,
            "chunk_count": len(range_chunks),
            "chunks_embedded_so_far": chunk_offset,
        })
        try:
            vs.create(range_chunks, tags=tags if tags else None)
        except Exception as e:
            doc_store.set_status(file_hash, "failed")
            _err(f"Embedding failed on pages {start}-{end}: {e}", "embedding_failed")

        if args.extract_figures:
            all_chunks_for_figures.extend(range_chunks)

        chunk_offset += len(range_chunks)
        total_chunks += len(range_chunks)

    doc_store.set_status(file_hash, "complete")
    doc_store.update_chunk_count(file_hash, total_chunks)

    # Store tags in manifest
    if tags:
        try:
            doc_store.update_tags(file_hash, tags)
        except Exception:
            pass  # Non-fatal

    # Extract figures if requested
    figure_count = 0
    if args.extract_figures and all_chunks_for_figures:
        _progress({"status": "extracting_figures"})
        figure_count = _extract_and_save_figures(all_chunks_for_figures, doc_store, file_hash)

    r2_tags: dict = {}
    if r2_mod.is_configured():
        r2_tags = _r2_upload_and_tag(args, doc_store, vs, file_hash, figure_count > 0)

    _ok({
        "status": "ingested",
        "collection": args.collection,
        "file_hash": file_hash,
        "document_name": first_dl_doc.name,
        "filename": first_dl_doc.origin.filename,
        "page_count": total_pages,
        "chunk_count": total_chunks,
        "figure_count": figure_count,
        "tags": tags or {},
        **r2_tags,
    })


def cmd_query(args: argparse.Namespace) -> None:
    # --file-hash may be specified multiple times; normalise to None / str / list
    hashes = args.file_hash or []
    if len(hashes) == 0:
        file_hash = None
    elif len(hashes) == 1:
        file_hash = hashes[0]
    else:
        file_hash = hashes

    tags = _parse_tags(getattr(args, "tag", None) or [])

    doc_store, vs = _stores(args)

    # Resolve --name to file_hash(es) when no explicit hash given
    if args.name and not hashes:
        file_hash = _resolve_name_filter(args, doc_store)

    try:
        results = vs.query(
            query_text=args.query,
            top_k=args.top_k,
            file_hash=file_hash,
            window=args.window,
            tags=tags if tags else None,
        )
    except Exception as e:
        _err(f"Query failed: {e}", "query_failed")

    _ok({
        "query": args.query,
        "top_k": args.top_k,
        "window": args.window,
        "tags": tags,
        "result_count": len(results),
        "results": [
            {
                "chunk": r.chunk.model_dump(),
                "context": [c.model_dump() for c in r.context],
            }
            for r in results
        ],
    })


def cmd_list(args: argparse.Namespace) -> None:
    tags = _parse_tags(getattr(args, "tag", None) or [])

    doc_store, vs = _stores(args)

    if args.source == "vector":
        try:
            docs = vs.list_documents()
        except Exception as e:
            _err(f"List failed: {e}", "list_failed")
        if args.name:
            needle = args.name.lower()
            docs = [d for d in docs if needle in d["document_name"].lower()]
        _ok({"count": len(docs), "source": "vector", "documents": docs})
        return

    try:
        records = doc_store.list()
    except Exception as e:
        _err(f"List failed: {e}", "list_failed")

    if args.name:
        needle = args.name.lower()
        records = [r for r in records if needle in r.document_name.lower() or needle in r.filename.lower()]

    # Filter by tags (all supplied tag keys must match; list values = OR within a key)
    if tags:
        filtered = []
        for r in records:
            match = True
            for k, v in tags.items():
                doc_val = r.tags.get(k)
                if isinstance(v, list):
                    if doc_val not in v:
                        match = False
                        break
                else:
                    if doc_val != v:
                        match = False
                        break
            if match:
                filtered.append(r)
        records = filtered

    _ok({
        "count": len(records),
        "source": "manifest",
        "documents": [r.model_dump() for r in records],
    })


def cmd_tag(args: argparse.Namespace) -> None:
    """Set or update tags on a document (manifest + vector store).

    Tags are merged into the existing tag set; existing tags not mentioned are
    preserved.  Provide tags as positional arguments in ``key=value`` format.
    Tag values are automatically normalized (lowercase, spaces → hyphens).

    Use ``--remove key`` to delete a tag key.
    """
    remove_keys = [k.strip().lower() for k in (args.remove or [])]
    tags = _parse_tags(args.tags or [], normalize=True)

    if not tags and not remove_keys:
        _err("No tags or --remove keys provided", "invalid_tag")

    doc_store, vs = _stores(args)

    # Resolve document
    file_hash = args.file_hash or _resolve_document(args.name, doc_store)

    # Remove tags first
    if remove_keys:
        try:
            doc_store.remove_tags(file_hash, remove_keys)
        except KeyError:
            _err(f"No document found with hash: {file_hash}", "not_found")
        except Exception as e:
            _err(f"Tag removal failed: {e}", "tag_failed")
        try:
            vs.remove_tag_metadata(file_hash, remove_keys)
        except Exception as e:
            _err(f"Vector store tag removal failed: {e}", "tag_failed")

    # Set/merge new tags
    if tags:
        # For ChromaDB metadata, list values (OR tags) are stored as comma-joined strings
        chroma_tags = {k: ",".join(v) if isinstance(v, list) else v for k, v in tags.items()}

        try:
            # Manifest only stores plain str values; store comma-joined for list values
            doc_store.update_tags(file_hash, chroma_tags)
        except KeyError:
            _err(f"No document found with hash: {file_hash}", "not_found")
        except Exception as e:
            _err(f"Tag update failed: {e}", "tag_failed")

        try:
            vs.set_tag_metadata(file_hash, chroma_tags)
        except Exception as e:
            _err(f"Vector store tag update failed: {e}", "tag_failed")

    _ok({"status": "tagged", "file_hash": file_hash, "tags_set": tags, "tags_removed": remove_keys})


def cmd_tags(args: argparse.Namespace) -> None:
    """Discover all tag keys and their distinct values across the collection."""
    doc_store, _vs = _stores(args)

    try:
        records = doc_store.list()
    except Exception as e:
        _err(f"Tags discovery failed: {e}", "list_failed")

    tags: Dict[str, list] = {}
    for r in records:
        for k, v in (r.tags or {}).items():
            if k not in tags:
                tags[k] = []
            if v not in tags[k]:
                tags[k].append(v)

    # Sort keys and values for stable output
    tags = {k: sorted(tags[k]) for k in sorted(tags)}

    _ok({"tags": tags})


def cmd_status(args: argparse.Namespace) -> None:
    doc_store, vs = _stores(args)

    try:
        records = doc_store.list()
    except Exception as e:
        _err(f"Status check failed: {e}", "status_failed")

    complete = [r for r in records if r.status == "complete"]
    incomplete = [r for r in records if r.status != "complete"]

    # Check which complete documents are missing from ChromaDB using per-doc existence
    # checks instead of a full collection scan. For small collections (typical case)
    # this is much faster than scanning all chunks.
    missing_in_chroma = []
    try:
        for r in complete:
            if not vs.has_chunks(r.file_hash):
                missing_in_chroma.append(r)
    except Exception as e:
        _err(f"Status check failed: {e}", "status_failed")

    # Use cached chunk counts from the manifest (written during ingest).
    # Documents ingested before this field was added will show 0.
    total_chunks = sum(r.chunk_count for r in complete)

    _ok({
        "collection": args.collection,
        "document_count": len(complete),
        "chunk_count": total_chunks,
        "incomplete_count": len(incomplete),
        "missing_in_chroma_count": len(missing_in_chroma),
        "healthy": len(incomplete) == 0 and len(missing_in_chroma) == 0,
    })


def cmd_repair(args: argparse.Namespace) -> None:
    doc_store, vs = _stores(args)

    try:
        incomplete = doc_store.list_incomplete()
        chroma_docs = vs.list_documents()
    except Exception as e:
        _err(f"Repair scan failed: {e}", "status_failed")

    chroma_hashes = {d["file_hash"] for d in chroma_docs}
    try:
        all_records = doc_store.list()
    except Exception as e:
        _err(f"Repair scan failed reading manifest: {e}", "status_failed")

    missing_in_chroma = [
        r for r in all_records
        if r.status == "complete" and r.file_hash not in chroma_hashes
    ]

    total_issues = len(incomplete) + len(missing_in_chroma)
    issues = {
        "incomplete": [r.model_dump() for r in incomplete],
        "missing_in_chroma": [r.model_dump() for r in missing_in_chroma],
    }

    if not args.fix:
        _ok({"issue_count": total_issues, **issues})
        return

    to_fix = {r.file_hash: r for r in incomplete + missing_in_chroma}
    fixed, failed = [], []
    for record in to_fix.values():
        try:
            pdf_path = doc_store.get_pdf_path(record.file_hash)
        except FileNotFoundError:
            failed.append({"file_hash": record.file_hash, "reason": "PDF not found in store"})
            continue
        try:
            converter = Converter()
            pdf_str = str(pdf_path)
            chunk_size = getattr(args, "chunk_size", 25)
            all_chunks: list = []
            chunk_offset = 0
            for _s, _e, doc in converter.convert_in_page_chunks(pdf_str, chunk_size):
                for chunk in Chunker(doc).chunk():
                    chunk.index = chunk_offset
                    chunk.id = f"{record.file_hash}_{chunk_offset}"
                    all_chunks.append(chunk)
                    chunk_offset += 1
            vs.update(all_chunks)
            doc_store.set_status(record.file_hash, "complete")
            fixed.append({"file_hash": record.file_hash, "document_name": record.document_name})
        except Exception as e:
            doc_store.set_status(record.file_hash, "failed")
            failed.append({"file_hash": record.file_hash, "reason": str(e)})

    _ok({"fixed": fixed, "failed": failed, **issues})


def _resolve_name_filter(args, doc_store):
    """Resolve --name to a file_hash filter for multi-match commands (query, ask).

    Returns None (no filter), a single hash string, or a list of hashes when
    multiple documents match the substring.
    """
    if not args.name:
        return None
    needle = args.name.lower()
    matched = [
        r.file_hash for r in doc_store.list()
        if needle in r.document_name.lower() or needle in r.filename.lower()
    ]
    if not matched:
        _err(f"No documents found matching name: {args.name}", "not_found")
    return matched[0] if len(matched) == 1 else matched


def _resolve_document(name: str, doc_store) -> str:
    """Resolve --name to exactly one file_hash for commands that require a single document.

    Exits with ``not_found`` if no documents match, or ``ambiguous_match`` if
    more than one document matches the substring.

    Args:
        name: Substring to match against document_name and filename (case-insensitive).
        doc_store: DocumentStore to search.

    Returns:
        The matching document's file_hash.
    """
    needle = name.lower()
    matched = [
        r for r in doc_store.list()
        if needle in r.document_name.lower() or needle in r.filename.lower()
    ]
    if not matched:
        _err(f"No documents found matching name: {name!r}", "not_found")
    if len(matched) > 1:
        _err(f"Multiple documents match '{name}' — use --file-hash", "ambiguous_match")
    return matched[0].file_hash


def cmd_ask(args: argparse.Namespace) -> None:
    doc_store, vs = _stores(args)

    # --file-hash takes precedence over --name when both are given
    hashes = getattr(args, "file_hash", None) or []
    if hashes:
        if len(hashes) == 1:
            file_hash = hashes[0]
        else:
            file_hash = hashes
    else:
        file_hash = _resolve_name_filter(args, doc_store)

    _progress({"status": "expanding_query", "question": args.question})
    try:
        keyphrases = expand_query(args.question)
    except Exception as e:
        _err(f"Query expansion failed: {e}", "query_expansion_failed")

    _progress({"status": "searching", "expanded_query": keyphrases})
    try:
        results = vs.query(
            query_text=keyphrases,
            top_k=args.top_k,
            file_hash=file_hash,
            window=args.window,
        )
    except Exception as e:
        _err(f"Query failed: {e}", "query_failed")

    if not results:
        _ok({
            "answer": "No relevant documents were found for your question.",
            "expanded_query": keyphrases,
            "sources": [],
        })
        return

    _progress({"status": "answering", "result_count": len(results)})
    try:
        answer = answer_question(args.question, results)
    except Exception as e:
        _err(f"Answer generation failed: {e}", "answer_failed")

    sources = []
    for r in results:
        sources.append({
            "document": r.chunk.document_name,
            "page": r.chunk.page_number,
            "file_hash": r.chunk.file_hash,
        })

    _ok({
        "answer": answer,
        "expanded_query": keyphrases,
        "sources": sources,
    })


def cmd_r2_upload(args: argparse.Namespace) -> None:
    """Upload an already-ingested document's PDF and figures to R2."""
    if not r2_mod.is_configured():
        _err("R2 not configured", "r2_not_configured")
    doc_store, vs = _stores(args)

    if getattr(args, "upload_all", False):
        records = [
            r for r in doc_store.list()
            if r.status == "complete" and "r2_pdf_key" not in (r.tags or {})
        ]
    else:
        file_hash = args.file_hash or _resolve_document(args.name, doc_store)
        records = [r for r in doc_store.list() if r.file_hash == file_hash]

    results = []
    for record in records:
        try:
            has_figures = bool(record.figure_count and record.figure_count > 0)
            new_tags = _r2_upload_and_tag(args, doc_store, vs, record.file_hash, has_figures)
            results.append({
                "file_hash": record.file_hash,
                "document_name": record.document_name,
                "status": "uploaded",
                "tags": new_tags,
            })
        except SystemExit:
            raise
        except Exception as e:
            results.append({
                "file_hash": record.file_hash,
                "document_name": record.document_name,
                "status": "failed",
                "error": str(e),
            })

    uploaded = len([r for r in results if r["status"] == "uploaded"])
    _ok({"uploaded": uploaded, "results": results})


def cmd_r2_url(args: argparse.Namespace) -> None:
    """Generate a presigned URL for a document PDF or figure stored in R2."""
    if not r2_mod.is_configured():
        _err("R2 not configured", "r2_not_configured")
    doc_store, _vs = _stores(args)

    file_hash = args.file_hash or _resolve_document(args.name, doc_store)
    records = {r.file_hash: r for r in doc_store.list()}
    record = records.get(file_hash)
    if not record:
        _err(f"Document not found: {file_hash}", "not_found")

    if args.figure:
        key = f"figures/{file_hash}/{args.figure}"
    else:
        key = (record.tags or {}).get("r2_pdf_key")
        if not key:
            _err("Document has no r2_pdf_key tag — run r2-upload first", "r2_not_uploaded")

    try:
        url = r2_mod.generate_presigned_url(key, expires_in=args.expires)
    except Exception as e:
        _err(f"Failed to generate presigned URL: {e}", "r2_url_failed")

    _ok({
        "file_hash": file_hash,
        "document_name": record.document_name,
        "key": key,
        "url": url,
        "expires_in": args.expires,
    })


def cmd_collections(args: argparse.Namespace) -> None:
    """List all collections in the data directory.

    A directory is considered a collection if it contains ``manifest.json``
    or a ``documents/`` subdirectory.
    """
    data_dir = Path(args.data_dir)
    if not data_dir.is_dir():
        _ok({"count": 0, "collections": []})
        return

    collections = []
    for child in sorted(data_dir.iterdir()):
        # Detect collections by presence of manifest.json or documents/ dir
        if child.is_dir() and (
            (child / "manifest.json").exists() or (child / "documents").is_dir()
        ):
            doc_store = DocumentStore(str(child))
            try:
                records = doc_store.list()
                doc_count = sum(1 for r in records if r.status == "complete")
            except Exception:
                doc_count = 0
            collections.append({
                "name": child.name,
                "document_count": doc_count,
                "path": str(child.relative_to(Path.cwd())) if child.is_relative_to(Path.cwd()) else str(child),
            })

    _ok({"count": len(collections), "collections": collections})


def cmd_figures(args: argparse.Namespace) -> None:
    doc_store, _vs = _stores(args)

    # Resolve target document
    file_hash = args.file_hash or _resolve_document(args.name, doc_store)

    figures = doc_store.get_figures(file_hash)

    if args.page:
        figures = [f for f in figures if f.page_number == args.page]

    if args.export_dir:
        export_path = Path(args.export_dir)
        export_path.mkdir(parents=True, exist_ok=True)
        exported = 0
        for fig in figures:
            try:
                src = doc_store.get_figure_image_path(file_hash, fig.id)
                dest = export_path / f"{fig.id}.png"
                import shutil
                shutil.copy2(src, dest)
                exported += 1
            except FileNotFoundError:
                pass
        _ok({"status": "exported", "file_hash": file_hash, "exported": exported, "target": str(export_path)})
        return

    _ok({
        "file_hash": file_hash,
        "figure_count": len(figures),
        "figures": [json.loads(f.model_dump_json()) for f in figures],
    })


def cmd_rename(args: argparse.Namespace) -> None:
    doc_store, vs = _stores(args)

    # Resolve by name if no hash given
    file_hash = args.file_hash or _resolve_document(args.name, doc_store)

    try:
        doc_store.rename(file_hash, args.new_name)
        vs.rename_document(file_hash, args.new_name)
    except KeyError:
        _err(f"No document found with hash: {file_hash}", "not_found")
    except Exception as e:
        _err(f"Rename failed: {e}", "rename_failed")

    _ok({"status": "renamed", "file_hash": file_hash, "new_name": args.new_name})


def cmd_delete(args: argparse.Namespace) -> None:
    doc_store, vs = _stores(args)

    try:
        doc_store.delete(args.file_hash)
        vs.delete(args.file_hash)
    except Exception as e:
        _err(f"Delete failed: {e}", "delete_failed")

    _ok({"status": "deleted", "file_hash": args.file_hash})


# ---------------------------------------------------------------------------
# Argument parser
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="vector",
        description="Vector document ingestion and semantic search CLI.",
    )
    parser.add_argument(
        "--data-dir", default=DEFAULT_DATA_DIR, metavar="DIR",
        help=f"Root directory for all collections (default: {DEFAULT_DATA_DIR}, env: VECTOR_DATA_DIR).",
    )
    parser.add_argument(
        "--collection", default=DEFAULT_COLLECTION, metavar="NAME",
        help=f"Collection to operate on (default: {DEFAULT_COLLECTION}). "
             "Data is stored under <data-dir>/<collection>/.",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # ingest
    p_ingest = sub.add_parser("ingest", help="Convert, store, chunk, and embed a PDF.")
    p_ingest.add_argument("file", help="Path to the PDF file.")
    p_ingest.add_argument(
        "--chunk-size", type=int, default=25, dest="chunk_size", metavar="N",
        help="Pages per conversion chunk (default: 25). "
             "Set to 0 to convert the whole document at once.",
    )
    p_ingest.add_argument(
        "--force", action="store_true",
        help="Delete and re-ingest the document even if it was already ingested.",
    )
    p_ingest.add_argument(
        "--extract-figures", action="store_true", dest="extract_figures",
        help="Extract figures (pictures and tables) as PNGs during ingestion.",
    )
    p_ingest.add_argument(
        "--tag", action="append", metavar="KEY=VALUE",
        help="Attach a tag to the document (repeatable). Format: key=value.",
    )
    p_ingest.add_argument(
        "--upload-r2", action="store_true", dest="upload_r2",
        help="Upload the source PDF and extracted figures to Cloudflare R2 after ingestion. "
             "Requires R2_ACCOUNT_ID, R2_BUCKET_NAME, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY "
             "in environment or .env.",
    )

    # query
    p_query = sub.add_parser("query", help="Semantic search over ingested documents.")
    p_query.add_argument("query", help="Natural-language query string.")
    p_query.add_argument("--top-k", type=int, default=5, metavar="N",
                         help="Maximum number of results (default: 5).")
    p_query.add_argument("--file-hash", action="append", metavar="HASH",
                         help="Restrict to one document. Repeat to search multiple documents.")
    p_query.add_argument("--name", metavar="SUBSTR",
                         help="Restrict to documents whose name contains this substring (case-insensitive). Resolved to file_hash automatically.")
    p_query.add_argument("--window", type=int, default=0, metavar="N",
                         help="Number of adjacent chunks to include around each hit (default: 0).")
    p_query.add_argument(
        "--tag", action="append", metavar="KEY=VALUE",
        help="Filter results to chunks tagged with this value (repeatable, AND logic). Format: key=value.",
    )

    # list
    p_list = sub.add_parser("list", help="List ingested documents.")
    p_list.add_argument("--name", metavar="SUBSTR",
                        help="Filter by document name or filename substring (case-insensitive).")
    p_list.add_argument("--source", choices=["manifest", "vector"], default="manifest",
                        help="Source to list from: 'manifest' (default) or 'vector' (ChromaDB).")
    p_list.add_argument(
        "--tag", action="append", metavar="KEY=VALUE",
        help="Filter to documents with matching tag(s) (repeatable, AND logic). Format: key=value.",
    )

    # status
    sub.add_parser("status", help="Show collection health: doc count, chunk count, and any issues.")

    # repair
    p_repair = sub.add_parser("repair", help="Find and optionally fix incomplete ingestions.")
    p_repair.add_argument("--fix", action="store_true",
                          help="Re-ingest each incomplete or missing document from its stored PDF.")

    # ask
    p_ask = sub.add_parser("ask", help="Ask a question: expand query, search, and answer.")
    p_ask.add_argument("question", help="Natural-language question.")
    p_ask.add_argument("--top-k", type=int, default=5, metavar="N",
                       help="Maximum number of search results (default: 5).")
    p_ask.add_argument("--window", type=int, default=0, metavar="N",
                       help="Adjacent chunks to include around each hit (default: 0).")
    p_ask.add_argument("--name", metavar="SUBSTR",
                       help="Restrict to documents whose name contains this substring.")
    p_ask.add_argument("--file-hash", action="append", metavar="HASH", dest="file_hash",
                       help="Restrict to a specific document by hash. Repeat for multiple.")

    # tag
    p_tag = sub.add_parser("tag", help="Set, update, or remove tags on a document.")
    p_tag_target = p_tag.add_mutually_exclusive_group(required=True)
    p_tag_target.add_argument("--file-hash", metavar="HASH", dest="file_hash",
                               help="Hash of the document to tag.")
    p_tag_target.add_argument("--name", metavar="SUBSTR",
                               help="Substring to match the document name (must be unambiguous).")
    p_tag.add_argument("tags", nargs="*", metavar="KEY=VALUE",
                       help="Tags to set in key=value format. Values are auto-normalized "
                            "(lowercase, spaces→hyphens). Existing tags not listed are preserved.")
    p_tag.add_argument("--remove", action="append", metavar="KEY", dest="remove",
                       help="Tag key to remove (repeatable). Example: --remove state --remove year.")

    # tags (discovery)
    sub.add_parser("tags", help="List all tag keys and their distinct values across the collection.")

    # figures
    p_figures = sub.add_parser("figures", help="List or export extracted figures for a document.")
    p_figures_target = p_figures.add_mutually_exclusive_group(required=True)
    p_figures_target.add_argument("--file-hash", metavar="HASH", dest="file_hash",
                                  help="Hash of the document.")
    p_figures_target.add_argument("--name", metavar="SUBSTR",
                                  help="Substring to match the document name.")
    p_figures.add_argument("--page", type=int, metavar="N",
                           help="Filter figures to a specific page number.")
    p_figures.add_argument("--export", metavar="DIR", dest="export_dir",
                           help="Export figure PNGs to a directory.")

    # rename
    p_rename = sub.add_parser("rename", help="Rename a document in the manifest and vector store.")
    p_rename_group = p_rename.add_mutually_exclusive_group(required=True)
    p_rename_group.add_argument("--file-hash", metavar="HASH", dest="file_hash",
                                help="Hash of the document to rename.")
    p_rename_group.add_argument("--name", metavar="SUBSTR",
                                help="Substring to match the document name (must be unambiguous).")
    p_rename.add_argument("new_name", metavar="NEW_NAME", help="New display name for the document.")

    # collections
    sub.add_parser("collections", help="List all collections in the data directory.")

    # delete
    p_delete = sub.add_parser("delete", help="Remove a document from all stores.")
    p_delete.add_argument("file_hash", help="The binary_hash of the document to remove.")

    # r2-upload
    p_r2upload = sub.add_parser(
        "r2-upload",
        help="Upload an already-ingested document PDF and figures to Cloudflare R2.",
    )
    p_r2upload_target = p_r2upload.add_mutually_exclusive_group()
    p_r2upload_target.add_argument("--file-hash", metavar="HASH", dest="file_hash",
                                   help="Hash of the document to upload.")
    p_r2upload_target.add_argument("--name", metavar="SUBSTR",
                                   help="Substring to match the document name.")
    p_r2upload.add_argument(
        "--all", action="store_true", dest="upload_all",
        help="Upload all complete documents that do not yet have an r2_pdf_key tag.",
    )

    # r2-url
    p_r2url = sub.add_parser(
        "r2-url",
        help="Generate a presigned URL for a document PDF or figure stored in R2.",
    )
    p_r2url_target = p_r2url.add_mutually_exclusive_group(required=True)
    p_r2url_target.add_argument("--file-hash", metavar="HASH", dest="file_hash",
                                help="Hash of the document.")
    p_r2url_target.add_argument("--name", metavar="SUBSTR",
                                help="Substring to match the document name.")
    p_r2url.add_argument("--figure", metavar="FILENAME",
                         help="Figure filename (e.g. pic_000.png) to generate URL for instead of the PDF.")
    p_r2url.add_argument("--expires", type=int, default=3600, metavar="SECONDS",
                         help="URL expiry in seconds (default: 3600).")

    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    dispatch = {
        "ingest": cmd_ingest,
        "query": cmd_query,
        "ask": cmd_ask,
        "list": cmd_list,
        "status": cmd_status,
        "repair": cmd_repair,
        "tag": cmd_tag,
        "tags": cmd_tags,
        "figures": cmd_figures,
        "rename": cmd_rename,
        "collections": cmd_collections,
        "delete": cmd_delete,
        "r2-upload": cmd_r2_upload,
        "r2-url": cmd_r2_url,
    }
    dispatch[args.command](args)


if __name__ == "__main__":
    main()