"""Command-line interface wiring for vectorize."""

from __future__ import annotations
from .validation import validate_config_sources
from .utils import (
    format_int,
    get_token_encoder,
    load_completion_state,
    summarize_collection,
    TOKEN_SAFETY_LIMIT,
)
from idxr.load_model_registry import load_model_registry
from idxr.models import ModelSpec
from .partitions import PARTITION_MANIFEST_VERSION, load_partition_manifest_entries
from .logging_config import setup_logging
from .indexing import (
    InvalidCollectionError,
    create_embedding_function,
    index_from_config,
)
from .e2e import E2ETestConfig, E2ETestRecorder
from .collection_strategies import (
    CollectionStrategy,
    FixedCollectionStrategy,
    PartitionCollectionStrategy,
)
from .configuration import ModelConfig, generate_stub_config, load_config
from .query_config import generate_query_config

import argparse
import json
import logging
import os
import random
import shutil
import textwrap
import threading
import time
from collections import deque
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Deque, Dict, List, Mapping, Optional, Sequence, Tuple

import chromadb

_OpenAIRateLimitError: Optional[type[BaseException]]
try:  # pragma: no cover - optional dependency probing
    # type: ignore[attr-defined]
    from openai import RateLimitError as _OpenAIRateLimitErrorClass
except ImportError:  # pragma: no cover - optional dependency probing
    _OpenAIRateLimitError = None
else:
    _OpenAIRateLimitError = _OpenAIRateLimitErrorClass

_RATE_LIMIT_ERROR_TYPES: Tuple[type[BaseException], ...] = (
    (_OpenAIRateLimitError,) if isinstance(_OpenAIRateLimitError, type) else ()
)


def _coerce_optional_int(value: Any) -> Optional[int]:
    """Best-effort conversion to int, returning None when not possible."""
    if isinstance(value, bool):
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, str):
        stripped = value.strip()
        if not stripped:
            return None
        try:
            return int(stripped)
        except ValueError:
            return None
    return None


def _iter_exception_chain(error: BaseException):
    """Yield an exception and its chained causes/contexts exactly once."""
    seen: set[int] = set()
    stack = [error]
    while stack:
        current = stack.pop()
        if not isinstance(current, BaseException):
            continue
        identity = id(current)
        if identity in seen:
            continue
        seen.add(identity)
        yield current
        cause = getattr(current, "__cause__", None)
        context = getattr(current, "__context__", None)
        if isinstance(cause, BaseException):
            stack.append(cause)
        if isinstance(context, BaseException):
            stack.append(context)


def _is_rate_limit_error(exc: BaseException) -> bool:
    """Return True when the exception chain includes an OpenAI rate limit error."""
    if not _RATE_LIMIT_ERROR_TYPES:
        return False
    return any(
        isinstance(candidate, _RATE_LIMIT_ERROR_TYPES)
        for candidate in _iter_exception_chain(exc)
    )


def _extract_retry_after_seconds(exc: BaseException) -> Optional[float]:
    """Extract the suggested retry delay (in seconds) if present."""
    for candidate in _iter_exception_chain(exc):
        retry_after = getattr(candidate, "retry_after", None)
        if isinstance(retry_after, (int, float)):
            return float(retry_after)
        error_payload = getattr(candidate, "error", None)
        if isinstance(error_payload, Mapping):
            direct = error_payload.get("retry_after")
            if isinstance(direct, (int, float)):
                return float(direct)
            retry_ms = error_payload.get("retry_after_ms")
            if isinstance(retry_ms, (int, float)):
                return float(retry_ms) / 1000.0
    return None


CLI_DESCRIPTION = textwrap.dedent(
    """
    Build and maintain a persistent ChromaDB index for SAP ECC 6.0 knowledge-base CSV exports.

    Provide --model <module:ATTRIBUTE> to point at the Pydantic ModelSpec registry you want to
    operate on (for the built-in ECC registry use
    kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY).

    You can ingest data directly from CSV mappings or from the manifest-driven partitions
    produced by prepare_datasets.py. A JSON configuration maps each Pydantic model name to its
    CSV definition, including the file path and any column renaming needed to match the schema.
    Use `init-config` to scaffold a full mapping, then trim it down to the models you actually
    plan to index. For incremental pipelines, let prepare_datasets.py accumulate partitions
    in a manifest and point `index` at that manifest to apply each “migration” in order.

    Examples
    --------
    # Generate a stub config to customize before indexing
    vectorize.py init-config \\
        --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --output data/vectorize_config.json

    # Index CSVs into a persistent collection
    vectorize.py index \\
        --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --config data/vectorize_config.json --collection ecc-std --persist-dir ./chroma

    # Index manifest-tracked partitions emitted by prepare_datasets.py
    vectorize.py index \\
        --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --partition-manifest build/partitions/manifest.json \\
        --partition-out-dir ./chroma_partitions --collection ecc-std --resume
    """
)

INDEX_DESCRIPTION = textwrap.dedent(
    """
    Load CSV data for the registered Pydantic models and upsert the content into a persistent
    ChromaDB collection. Supply either a JSON config that lists the CSV sources to ingest or a
    manifest of partitions produced by prepare_datasets.py. Model entries without a `path` are
    skipped with a warning. Optional `columns` mappings rename headers before validation.

    All configured CSVs are streamed once up front to confirm they exist and pass Pydantic
    validation before any data is upserted. When used with `--resume`, the command will reuse
    completion metadata from a JSON state file so that models already indexed (and whose source
    CSVs have not changed) are skipped entirely. Alternatively, provide `--partition-manifest`
    (or legacy `--partition-dir`) together with `--partition-out-dir` to index partition
    directories in manifest order, keeping independent persistence and resume state per
    partition.
    """
)

DISPLAY_DESCRIPTION = textwrap.dedent(
    """
    Inspect an existing ChromaDB collection and print high-level statistics,
    including total document count and per-model document totals based on stored metadata.

    Examples
    --------
    vectorize.py display --collection ecc-std --persist-dir ./chroma --top 10
    """
)

INIT_DESCRIPTION = textwrap.dedent(
    """
    Write a JSON file containing every supported Pydantic model name mapped to an object
    with `path` and `columns` keys. The generated `columns` map lists the model's
    field names so you can fill in the actual CSV headers before running `index`.
    Remove entries for models you do not plan to index.
    """
)

DROP_DESCRIPTION = textwrap.dedent(
    """
    Delete indexed documents for specific models using a drop configuration generated by
    prepare_datasets.py plan-drop/apply-drop. Supply --apply to execute the deletion, otherwise
    a dry-run summary is printed.

    Example
    -------
    vectorize.py apply-drop --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --config configs/drop/gdpr_2024-07.json --collection ecc-std
    """
)

STATUS_DESCRIPTION = textwrap.dedent(
    """
    Display the current indexing progress by analyzing resume state files and error reports.
    Shows which partitions are started, finished, or have errors.

    Examples
    --------
    # Basic usage (shows progress from resume state files and errors)
    vectorize.py status --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --partition-out-dir build/vector

    # With config directory (shows which configured models haven't started)
    vectorize.py status --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --partition-out-dir build/vector \\
        --partition-dir build/partitions
    """
)


class VectorizeCLI:
    """Command-line interface entry point."""

    def __init__(self) -> None:
        self.parser = argparse.ArgumentParser(
            prog="vectorize.py",
            description=CLI_DESCRIPTION,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )
        self._add_subcommands()

    def _add_subcommands(self) -> None:
        common = argparse.ArgumentParser(
            add_help=False,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )
        common.add_argument(
            "--model",
            required=True,
            help=(
                "Python import string pointing to the model registry "
                "(format: package.module:REGISTRY_NAME)."
            ),
        )
        common.add_argument(
            "--log-level",
            default="INFO",
            help="Logging verbosity (DEBUG, INFO, WARNING, ERROR). Defaults to INFO.",
        )
        common.add_argument(
            "--log-file",
            type=Path,
            help=(
                "Path to log file. Enables file logging with automatic rotation. "
                "If not specified, logs only to console."
            ),
        )
        common.add_argument(
            "--log-max-bytes",
            type=int,
            default=100 * 1024 * 1024,  # 100 MB
            help=(
                "Maximum size of a single log file before rotation in bytes. "
                "Default: 104857600 (100 MB)."
            ),
        )
        common.add_argument(
            "--log-backup-count",
            type=int,
            default=10,
            help=(
                "Number of rotated log files to keep. "
                "Default: 10 (for 1GB total with 100MB files)."
            ),
        )
        common.add_argument(
            "--log-no-console",
            action="store_true",
            help="Disable console logging (only log to file if --log-file is specified).",
        )

        subparsers = self.parser.add_subparsers(dest="command", required=True)

        index_parser = subparsers.add_parser(
            "index",
            parents=[common],
            help="Index CSV data into a persistent ChromaDB collection.",
            description=INDEX_DESCRIPTION,
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog=textwrap.dedent(
                """
                Metadata tip: each document carries a `model_name` field,
                so you can filter search results with `where={"model_name": "Table"}`
                when querying ChromaDB.

                Examples
                --------
                # Fresh index run
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --config data/vectorize_config.json \\
                    --collection ecc-std --persist-dir ./chroma

                # Resume using the default state file
                # (<persist-dir>/<collection>_resume_state.json)
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --config data/vectorize_config.json \\
                    --collection ecc-std \\
                    --persist-dir ./chroma --resume

                # Index manifest-driven partitions
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --partition-manifest build/partitions/manifest.json \\
                    --partition-out-dir ./chroma_partitions --collection ecc-std

                # Smoke-test a manifest run by sampling 50 rows per CSV
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --partition-manifest build/partitions/manifest.json \\
                    --partition-out-dir ./chroma_partitions \\
                    --collection ecc-std \\
                    --e2e-test-run --e2e-sample-size 50 \\
                    --e2e-output build/vector/e2e_samples.json

                # Run partitions against Chroma Cloud
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --partition-manifest build/partitions/manifest.json \\
                    --partition-out-dir ./chroma_partitions \\
                    --client-type cloud \\
                    --chroma-api-token "$CHROMA_TOKEN" \\
                    --collection ecc-std

                # Sample 25 rows per CSV when targeting Chroma Cloud
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --partition-manifest build/partitions/manifest.json \\
                    --partition-out-dir ./chroma_partitions \\
                    --client-type cloud \\
                    --chroma-api-token "$CHROMA_TOKEN" \\
                    --collection ecc-std \\
                    --e2e-test-run --e2e-sample-size 25 \\
                    --e2e-output build/vector/e2e_cloud_samples.json

                Logging Examples
                ----------------
                # Enable file logging with default rotation (100MB files, 10 backups)
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --partition-manifest build/partitions/manifest.json \\
                    --partition-out-dir build/vector \\
                    --collection ecc-std \\
                    --log-file logs/vectorize.log

                # Production: large files, no console output (for background jobs)
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --partition-manifest build/partitions/manifest.json \\
                    --partition-out-dir build/vector \\
                    --collection ecc-std \\
                    --log-file /var/log/vectorize/production.log \\
                    --log-max-bytes 524288000 \\
                    --log-backup-count 20 \\
                    --log-no-console

                # Debug mode with verbose console output only
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --config data/config.json \\
                    --collection ecc-std \\
                    --persist-dir ./chroma \\
                    --log-level DEBUG

                # Both console and file logging with custom rotation
                vectorize.py index \\
                    --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
                    --partition-manifest build/partitions/manifest.json \\
                    --partition-out-dir build/vector \\
                    --collection ecc-std \\
                    --log-file logs/index_$(date +%%Y%%m%%d_%%H%%M%%S).log \\
                    --log-max-bytes 209715200 \\
                    --log-backup-count 5 \\
                    --log-level INFO
                """
            ),
        )
        index_parser.add_argument(
            "--config",
            type=Path,
            help="JSON configuration mapping model names to CSV files.",
        )
        index_parser.add_argument(
            "--collection",
            help=(
                "Collection name or prefix, depending on the active strategy. "
                "Required for the persistent client; optional for partition-aware "
                "HTTP runs where each partition becomes its own collection."
            ),
        )
        index_parser.add_argument(
            "--persist-dir",
            type=Path,
            help=(
                "Directory where the ChromaDB persistent client stores data. "
                "Required when --client-type=persistent (default)."
            ),
        )
        index_parser.add_argument(
            "--batch-size",
            type=int,
            default=128,
            help=(
                "Maximum number of documents to send per embedding request "
                f"(defaults to 128; hard limit {format_int(TOKEN_SAFETY_LIMIT)} tokens)."
            ),
        )
        index_parser.add_argument(
            "--resume",
            action="store_true",
            help=(
                "Resume indexing by skipping rows that already exist in the target collection."
                " Combined with the resume state file, models whose source CSVs are unchanged"
                " are skipped entirely."
            ),
        )
        index_parser.add_argument(
            "--embedding-model",
            default="text-embedding-3-small",
            help=(
                "OpenAI embedding model identifier (e.g. text-embedding-3-small, "
                "text-embedding-3-large)."
            ),
        )
        index_parser.add_argument(
            "--openai-api-key",
            help="OpenAI API key (otherwise the OPENAI_API_KEY environment variable is used).",
        )
        index_parser.add_argument(
            "--resume-state-file",
            type=Path,
            help=(
                "Path to the JSON file that tracks per-model completion metadata when"
                " resuming. Defaults to <persist-dir>/<collection>_resume_state.json; the file"
                " is created automatically after each model completes."
            ),
        )
        index_parser.add_argument(
            "--skip-validation",
            action="store_true",
            help="Skip CSV validation step before indexing (use with caution).",
        )
        index_parser.add_argument(
            "--truncation-strategy",
            choices=["end", "start", "middle_out", "sentences", "auto"],
            default="auto",
            help=(
                "Default truncation strategy for documents exceeding token limits. "
                "Choices: 'end' (keep beginning), 'start' (keep end), "
                "'middle_out' (preserve both ends), 'sentences' (complete sentences), "
                "'auto' (automatically detect best strategy). "
                "Can be overridden per-model in config. Default: auto."
            ),
        )
        index_parser.add_argument(
            "--token-limit",
            type=int,
            default=TOKEN_SAFETY_LIMIT,
            help=(
                "Maximum tokens per batch request sent to embedding API. "
                "Total token limit for all documents in a single batch. "
                f"Default: {TOKEN_SAFETY_LIMIT}."
            ),
        )
        index_parser.add_argument(
            "--embedding-token-limit",
            type=int,
            default=8191,
            help=(
                "Maximum tokens per individual document for embedding. "
                "Documents exceeding this limit will be truncated. "
                "Default: 8191 (OpenAI embedding model limit)."
            ),
        )
        index_parser.add_argument(
            "--e2e-test-run",
            action="store_true",
            help=(
                "Enable end-to-end sampling mode: index a random subset of rows from each CSV "
                "and emit an audit log of sampled entries."
            ),
        )
        index_parser.add_argument(
            "--e2e-sample-size",
            type=int,
            default=100,
            help="Number of rows to sample per CSV when "
            "--e2e-test-run is enabled (default: 100).",
        )
        index_parser.add_argument(
            "--e2e-output",
            type=Path,
            default=Path("vectorize_e2e_samples.json"),
            help=(
                "Path to the audit file that records sampled rows during --e2e-test-run "
                "(defaults to ./vectorize_e2e_samples.json)."
            ),
        )
        index_parser.add_argument(
            "--e2e-random-seed",
            type=int,
            help="Optional random seed for reproducible --e2e-test-run sampling.",
        )
        index_parser.add_argument(
            "--client-type",
            choices=("persistent", "http", "cloud"),
            default="persistent",
            help="Chroma client implementation to use (default: persistent).",
        )
        index_parser.add_argument(
            "--chroma-server-host",
            help="Host for the HTTP Chroma client (required when --client-type=http).",
        )
        index_parser.add_argument(
            "--chroma-server-port",
            type=int,
            help="Port for the HTTP Chroma client (defaults to server default).",
        )
        index_parser.add_argument(
            "--chroma-server-ssl",
            action="store_true",
            help="Enable TLS/SSL when using the HTTP Chroma client.",
        )
        index_parser.add_argument(
            "--chroma-api-token",
            help="API token for the HTTP Chroma client "
            "(uses Authorization: Bearer ... header).",
        )
        index_parser.add_argument(
            "--chroma-tenant",
            help="Tenant identifier header for Chroma Cloud (sets X-Chroma-Tenant).",
        )
        index_parser.add_argument(
            "--chroma-database",
            help="Database identifier header for Chroma Cloud (sets X-Chroma-Database).",
        )
        index_parser.add_argument(
            "--chroma-cloud-tenant",
            help="Tenant identifier when using the Cloud client "
            "(equivalent to --chroma-tenant).",
        )
        index_parser.add_argument(
            "--chroma-cloud-database",
            help="Database identifier when using the Cloud client "
            "(equivalent to --chroma-database).",
        )
        index_parser.add_argument(
            "--chroma-cloud-host",
            help="Override the Chroma Cloud host (defaults to api.trychroma.com).",
        )
        index_parser.add_argument(
            "--chroma-cloud-port",
            type=int,
            help="Override the Chroma Cloud port (defaults to 443).",
        )
        index_parser.add_argument(
            "--partition-dir",
            type=Path,
            help=(
                "Directory containing partition subdirectories produced by prepare_datasets.py."
                " Each partition must include a vectorize config file."
            ),
        )
        index_parser.add_argument(
            "--partition-manifest",
            type=Path,
            help=(
                f"Path to a manifest.json produced by prepare_datasets.py "
                f"(version {PARTITION_MANIFEST_VERSION}); partitions are "
                "indexed in manifest order when supplied."
            ),
        )
        index_parser.add_argument(
            "--partition-out-dir",
            type=Path,
            help=(
                "Root directory where per-partition Chroma persistence data should be written."
                " Required when partition indexing is used."
            ),
        )
        index_parser.add_argument(
            "--parallel-partitions",
            type=int,
            default=1,
            help=(
                "Number of partitions to index concurrently when partition indexing is enabled."
                " Defaults to 1 (no parallelism)."
            ),
        )
        index_parser.add_argument(
            "--partition-config-name",
            default="vectorize_config.json",
            help=(
                "File name of the vectorize configuration within each partition directory "
                "(defaults to vectorize_config.json)."
            ),
        )
        index_parser.add_argument(
            "--delete-stale",
            action="store_true",
            help=(
                "When used with --partition-manifest, delete stale partitions from the target "
                "Chroma store after replacement partitions are indexed."
            ),
        )
        index_parser.set_defaults(func=self.handle_index)

        display_parser = subparsers.add_parser(
            "display",
            parents=[common],
            help="Display summary statistics for a ChromaDB collection.",
            description=DISPLAY_DESCRIPTION,
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog=textwrap.dedent(
                """
    Examples
    --------
    # Show top 10 models for a collection stored in ./chroma
    vectorize.py display \\
        --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --collection ecc-std --persist-dir ./chroma --top 10

    # Scan the entire collection with a larger chunk size
    vectorize.py display \\
        --model kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY \\
        --collection ecc-std --persist-dir ./chroma --chunk-size 20000 --top 0
                """
            ),
        )
        display_parser.add_argument(
            "--collection",
            required=True,
            help="Name of the ChromaDB collection to inspect.",
        )
        display_parser.add_argument(
            "--persist-dir",
            required=True,
            type=Path,
            help="Directory where the ChromaDB persistent client stores data.",
        )
        display_parser.add_argument(
            "--chunk-size",
            type=int,
            default=10000,
            help="Number of documents to read per batch when computing per-model counts.",
        )
        display_parser.add_argument(
            "--top",
            type=int,
            default=20,
            help="Number of per-model rows to display (set to 0 to show all).",
        )
        display_parser.set_defaults(func=self.handle_display)

        init_parser = subparsers.add_parser(
            "init-config",
            parents=[common],
            help="Generate a stub JSON configuration file.",
            description=INIT_DESCRIPTION,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )
        init_parser.add_argument(
            "--output",
            type=Path,
            default=Path("vectorize_config.json"),
            help=(
                "Destination for the generated config file."
                " Defaults to ./vectorize_config.json; directories are created as needed."
            ),
        )
        init_parser.add_argument(
            "--force",
            action="store_true",
            help="Overwrite the output file if it already exists.",
        )
        init_parser.set_defaults(func=self.handle_init_config)

        drop_parser = subparsers.add_parser(
            "apply-drop",
            parents=[common],
            help="Remove documents for specific models using a drop configuration.",
            description=DROP_DESCRIPTION,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )
        drop_parser.add_argument(
            "--config",
            required=True,
            type=Path,
            help="Drop configuration (JSON) listing the models/partitions to delete.",
        )
        drop_parser.add_argument(
            "--collection",
            required=True,
            help="Name of the ChromaDB collection to modify.",
        )
        drop_parser.add_argument(
            "--client-type",
            choices=("persistent", "http"),
            default="persistent",
            help="Chroma client implementation to use (default: persistent).",
        )
        drop_parser.add_argument(
            "--persist-dir",
            type=Path,
            help="Path to the persistent client directory "
            "(required for --client-type=persistent).",
        )
        drop_parser.add_argument(
            "--chroma-server-host",
            help="Host for the HTTP Chroma client (required when --client-type=http).",
        )
        drop_parser.add_argument(
            "--chroma-server-port",
            type=int,
            help="Port for the HTTP Chroma client (defaults to server default).",
        )
        drop_parser.add_argument(
            "--chroma-server-ssl",
            action="store_true",
            help="Enable TLS/SSL when using the HTTP Chroma client.",
        )
        drop_parser.add_argument(
            "--chroma-api-token",
            help="API token for the HTTP Chroma client (Authorization: Bearer ...).",
        )
        drop_parser.add_argument(
            "--chroma-tenant",
            help="Tenant identifier header for Chroma Cloud (sets X-Chroma-Tenant).",
        )
        drop_parser.add_argument(
            "--chroma-database",
            help="Database identifier header for Chroma Cloud (sets X-Chroma-Database).",
        )
        drop_parser.add_argument(
            "--partition-manifest",
            type=Path,
            help="Optional manifest to mark partitions as deleted after Chroma cleanup.",
        )
        drop_parser.add_argument(
            "--apply",
            action="store_true",
            help="Execute the deletion. Without this flag a dry-run summary is shown.",
        )
        drop_parser.add_argument(
            "--performed-by",
            help="Optional identifier recorded when updating the manifest.",
        )
        drop_parser.set_defaults(func=self.handle_drop_models)

        status_parser = subparsers.add_parser(
            "status",
            parents=[common],
            help="Display current indexing progress.",
            description=STATUS_DESCRIPTION,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )
        status_parser.add_argument(
            "--partition-out-dir",
            type=Path,
            required=True,
            help="Root directory where per-partition Chroma persistence data is written.",
        )
        status_parser.add_argument(
            "--partition-dir",
            type=Path,
            help=(
                "Directory containing partition subdirectories with "
                "vectorize_config.json files. "
                "Used to detect which models are configured but not yet started."
            ),
        )
        status_parser.add_argument(
            "--partition-config-name",
            default="vectorize_config.json",
            help=(
                "File name of the vectorize configuration within each partition directory "
                "(defaults to vectorize_config.json)."
            ),
        )
        status_parser.set_defaults(func=self.handle_status)

        # Generate query config command
        query_config_parser = subparsers.add_parser(
            "generate-query-config",
            parents=[common],
            help="Generate query configuration from partition resume states.",
            description=(
                "Generate a query_config.json file by scanning resume state files "
                "in partition output directories. This config maps model names to "
                "their collections for efficient multi-collection querying."
            ),
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )
        query_config_parser.add_argument(
            "--partition-out-dir",
            type=Path,
            required=True,
            help="Root directory containing partition subdirectories with resume state files.",
        )
        query_config_parser.add_argument(
            "--output",
            type=Path,
            required=True,
            help="Path where the query_config.json file will be written.",
        )
        query_config_parser.add_argument(
            "--collection-prefix",
            help="Optional collection name prefix used during indexing.",
        )
        query_config_parser.add_argument(
            "--generate-environment",
            action="store_true",
            help="Generate environment.json alongside query config with indexing settings.",
        )
        query_config_parser.add_argument(
            "--environment-output",
            type=Path,
            help="Path for environment.json (default: <partition-out-dir>/environment.json).",
        )
        query_config_parser.add_argument(
            "--embedding-model",
            default="text-embedding-3-small",
            help="Embedding model used during indexing (default: text-embedding-3-small).",
        )
        query_config_parser.add_argument(
            "--model-registry-target",
            help="Import path to model registry "
            "(e.g., 'kb.std.ecc_6_0_ehp_7.registry:MODEL_REGISTRY'). "
            "Required with --generate-environment.",
        )
        query_config_parser.add_argument(
            "--discriminator-field",
            default="model",
            help="Metadata field used as discriminator for model routing (default: model).",
        )
        query_config_parser.add_argument(
            "--chroma-client-type",
            choices=["persistent", "cloud", "http", "memory"],
            default="persistent",
            help="ChromaDB client type used during indexing (default: persistent).",
        )
        query_config_parser.add_argument(
            "--cloud-api-key",
            help="ChromaDB cloud API key. Use 'env:VAR_NAME' format to reference "
            "environment variable. Required for cloud client.",
        )
        query_config_parser.add_argument(
            "--cloud-tenant",
            help="ChromaDB cloud tenant ID. Required for cloud client.",
        )
        query_config_parser.add_argument(
            "--cloud-database",
            help="ChromaDB cloud database name. Required for cloud client.",
        )
        query_config_parser.add_argument(
            "--http-host",
            help="HTTP client host. Required for http client.",
        )
        query_config_parser.add_argument(
            "--http-port",
            type=int,
            help="HTTP client port. Required for http client.",
        )
        query_config_parser.add_argument(
            "--http-ssl",
            action="store_true",
            help="Use SSL for HTTP client.",
        )
        query_config_parser.set_defaults(func=self.handle_generate_query_config)

    def run(self, argv: Sequence[str]) -> int:
        args = self.parser.parse_args(argv)

        # Setup logging with rotation support
        setup_logging(
            log_level=str(args.log_level),
            log_file=getattr(args, "log_file", None),
            max_bytes=getattr(args, "log_max_bytes", 100 * 1024 * 1024),
            backup_count=getattr(args, "log_backup_count", 10),
            console_output=not getattr(args, "log_no_console", False),
        )

        try:
            args.model_registry = load_model_registry(str(args.model))
        except (TypeError, ValueError) as exc:
            logging.error("Failed to load model registry %s: %s", args.model, exc)
            return 1
        return args.func(args)

    @staticmethod
    def _resolve_http_settings(
        args: argparse.Namespace,
    ) -> Tuple[Optional[str], Optional[int], bool, Dict[str, str]]:
        host = args.chroma_server_host or os.getenv("CHROMA_SERVER_HOST")

        port: Optional[int] = args.chroma_server_port
        if port is None:
            port_env = os.getenv("CHROMA_SERVER_PORT")
            if port_env:
                try:
                    port = int(port_env)
                except ValueError:
                    logging.warning(
                        "Ignoring invalid CHROMA_SERVER_PORT value '%s'", port_env
                    )

        ssl_flag = bool(args.chroma_server_ssl)
        if not ssl_flag:
            ssl_env = os.getenv("CHROMA_SERVER_SSL")
            if ssl_env:
                ssl_flag = ssl_env.strip().lower() in {"1", "true", "yes", "on"}

        headers: Dict[str, str] = {}
        token = args.chroma_api_token or os.getenv("CHROMA_API_TOKEN")
        if token:
            headers["Authorization"] = f"Bearer {token}"
        tenant = args.chroma_tenant or os.getenv("CHROMA_TENANT")
        if tenant:
            headers["X-Chroma-Tenant"] = tenant
        database = args.chroma_database or os.getenv("CHROMA_DATABASE")
        if database:
            headers["X-Chroma-Database"] = database

        return host, port, ssl_flag, headers

    @staticmethod
    def _resolve_cloud_settings(
        args: argparse.Namespace,
    ) -> Tuple[Optional[str], Optional[str], Optional[str], str, int, bool]:
        tenant = (
            args.chroma_cloud_tenant
            or args.chroma_tenant
            or os.getenv("CHROMA_CLOUD_TENANT")
            or os.getenv("CHROMA_TENANT")
        )
        database = (
            args.chroma_cloud_database
            or args.chroma_database
            or os.getenv("CHROMA_CLOUD_DATABASE")
            or os.getenv("CHROMA_DATABASE")
        )
        api_key = args.chroma_api_token or os.getenv("CHROMA_API_TOKEN")
        host = (
            args.chroma_cloud_host
            or os.getenv("CHROMA_CLOUD_HOST")
            or "api.trychroma.com"
        )
        port_value = args.chroma_cloud_port or os.getenv("CHROMA_CLOUD_PORT")
        port = 443
        if port_value is not None:
            try:
                port = int(port_value)
            except (TypeError, ValueError):
                logging.warning(
                    "Ignoring invalid Chroma Cloud port value '%s'", port_value
                )
        ssl_flag = True
        return tenant, database, api_key, host, port, ssl_flag

    def _build_chroma_client(
        self,
        *,
        client_type: str,
        persist_dir: Optional[Path],
        http_host: Optional[str],
        http_port: Optional[int],
        http_ssl: bool,
        http_headers: Mapping[str, str],
        cloud_tenant: Optional[str] = None,
        cloud_database: Optional[str] = None,
        cloud_api_key: Optional[str] = None,
        cloud_host: str = "api.trychroma.com",
        cloud_port: int = 443,
        cloud_ssl: bool = True,
    ):
        """Create the ChromaDB client used for persistence."""
        if client_type == "persistent":
            if persist_dir is None:
                raise ValueError(
                    "persist_dir is required when using the persistent Chroma client."
                )
            return chromadb.PersistentClient(path=str(persist_dir))
        if client_type == "http":
            if not http_host:
                raise ValueError(
                    "--chroma-server-host must be provided when using the HTTP Chroma client."
                )
            headers = dict(http_headers) if http_headers else {}
            client_kwargs: Dict[str, Any] = {
                "host": http_host,
                "ssl": http_ssl,
                "headers": headers or None,
            }
            if http_port is not None:
                client_kwargs["port"] = int(http_port)
            return chromadb.HttpClient(**client_kwargs)
        if client_type == "cloud":
            if cloud_api_key is None:
                raise ValueError(
                    "--chroma-api-token (or CHROMA_API_TOKEN) "
                    "is required when --client-type=cloud."
                )
            return chromadb.CloudClient(
                tenant=cloud_tenant,
                database=cloud_database,
                api_key=cloud_api_key,
                cloud_host=cloud_host,
                cloud_port=int(cloud_port),
                enable_ssl=bool(cloud_ssl),
            )
        raise ValueError(f"Unsupported client type: {client_type}")

    def _create_collection(
        self,
        *,
        client_type: str,
        persist_dir: Optional[Path],
        collection_name: str,
        embedding_function: chromadb.EmbeddingFunction,
        http_host: Optional[str],
        http_port: Optional[int],
        http_ssl: bool,
        http_headers: Mapping[str, str],
        cloud_tenant: Optional[str],
        cloud_database: Optional[str],
        cloud_api_key: Optional[str],
        cloud_host: str,
        cloud_port: int,
        cloud_ssl: bool,
    ):
        """Create or retrieve the target collection."""
        client = self._build_chroma_client(
            client_type=client_type,
            persist_dir=persist_dir if client_type == "persistent" else None,
            http_host=http_host,
            http_port=http_port,
            http_ssl=http_ssl,
            http_headers=http_headers,
            cloud_tenant=cloud_tenant,
            cloud_database=cloud_database,
            cloud_api_key=cloud_api_key,
            cloud_host=cloud_host,
            cloud_port=cloud_port,
            cloud_ssl=cloud_ssl,
        )

        metadata: Dict[str, Any] = {}
        model_name_attr = getattr(embedding_function, "model_name", None)
        if isinstance(model_name_attr, str) and model_name_attr:
            metadata["embedding_model"] = model_name_attr

        if client_type == "persistent":
            collection = client.get_or_create_collection(
                name=collection_name,
                embedding_function=embedding_function,
                metadata=metadata or None,
            )
        elif client_type == "http":
            collection = client.get_or_create_collection(
                name=collection_name,
                embedding_function=embedding_function,
                metadata=metadata or None,
            )
        else:  # cloud client
            collection = client.get_or_create_collection(
                name=collection_name,
                embedding_function=embedding_function,
                metadata=metadata or None,
            )
        if metadata:
            existing_meta = getattr(collection, "metadata", None)
            if not isinstance(existing_meta, Mapping):
                existing_meta = None
            if existing_meta != metadata:
                try:
                    collection.modify(metadata=metadata)
                except Exception as exc:  # pragma: no cover - defensive
                    logging.debug(
                        "Unable to update collection metadata for %s: %s",
                        collection_name,
                        exc,
                    )
        return collection

    def _run_index_for_config(
        self,
        *,
        config_path: Path,
        collection_name: str,
        persist_dir: Path,
        batch_size: int,
        embedding_model: str,
        resume: bool,
        skip_validation: bool,
        embedding_function,
        encoder,
        resume_state_file: Optional[Path],
        client_type: str,
        http_host: Optional[str],
        http_port: Optional[int],
        http_ssl: bool,
        http_headers: Mapping[str, str],
        cloud_tenant: Optional[str],
        cloud_database: Optional[str],
        cloud_api_key: Optional[str],
        cloud_host: str,
        cloud_port: int,
        cloud_ssl: bool,
        model_registry: Mapping[str, ModelSpec],
        truncation_strategy: str = "auto",
        token_limit: int = TOKEN_SAFETY_LIMIT,
        embedding_token_limit: int = 8191,
        extra_metadata: Optional[Mapping[str, Any]] = None,
        model_metadata: Optional[Mapping[str, Mapping[str, Any]]] = None,
        e2e_config: Optional[E2ETestConfig] = None,
    ) -> int:
        try:
            config = load_config(config_path, model_registry)
        except (OSError, ValueError, json.JSONDecodeError, KeyError) as exc:
            logging.error("Failed to load config %s: %s", config_path, exc)
            return 1

        configured_models = len(config)
        active_models = sum(
            1 for cfg in config.values() if isinstance(cfg, ModelConfig) and cfg.path
        )
        logging.info(
            "Starting index run for %s (batch=%s, embedding=%s, resume=%s, token-limit=%s)",
            config_path,
            format_int(batch_size),
            embedding_model,
            "yes" if resume else "no",
            format_int(token_limit),
        )
        logging.info(
            "Loaded configuration for %d models (%d with CSV paths)",
            configured_models,
            active_models,
        )

        resume_state_path: Optional[Path] = (
            resume_state_file
            if resume_state_file is not None
            else persist_dir / f"{collection_name}_resume_state.json"
        )
        completion_state: Dict[str, Any] = (
            load_completion_state(resume_state_path) if resume else {}
        )
        if resume:
            completed_models = sum(
                1
                for entry in completion_state.values()
                if isinstance(entry, Mapping) and entry.get("complete")
            )
            logging.info(
                "Loaded completion metadata from %s (%d model(s) marked complete).",
                resume_state_path,
                completed_models,
            )

        if skip_validation:
            logging.warning(
                "Skipping CSV validation per --skip-validation flag for %s.",
                config_path,
            )
        else:
            logging.info("Validating configured CSV sources before indexing...")
            if not validate_config_sources(config, model_registry=model_registry):
                logging.error(
                    "Aborting indexing because one or more CSVs failed validation."
                )
                return 1
            logging.info("Validation complete; beginning data ingestion.")

        collection = self._create_collection(
            client_type=client_type,
            persist_dir=persist_dir if client_type == "persistent" else None,
            collection_name=collection_name,
            embedding_function=embedding_function,
            http_host=http_host,
            http_port=http_port,
            http_ssl=http_ssl,
            http_headers=http_headers,
            cloud_tenant=cloud_tenant,
            cloud_database=cloud_database,
            cloud_api_key=cloud_api_key,
            cloud_host=cloud_host,
            cloud_port=cloud_port,
            cloud_ssl=cloud_ssl,
        )

        counts = index_from_config(
            collection,
            config,
            batch_size=batch_size,
            model_registry=model_registry,
            resume=resume,
            encoder=encoder,
            token_limit=token_limit,
            embedding_token_limit=embedding_token_limit,
            completion_state=completion_state,
            completion_state_path=resume_state_path,
            extra_metadata=extra_metadata,
            model_metadata=model_metadata,
            e2e_config=e2e_config,
            error_report_dir=persist_dir,
            collection_name=collection_name,
            default_truncation_strategy=truncation_strategy,
        )
        total = sum(counts.values())
        logging.info("Indexed %d documents across %d models", total, len(counts))
        if resume_state_path and resume_state_path.exists():
            logging.info("Completion metadata stored at %s", resume_state_path)
        return 0

    def _initialize_embedding_resources(self, embedding_model: str, api_key: str):
        encoder = get_token_encoder(embedding_model)
        embedding_function = create_embedding_function(embedding_model, api_key=api_key)
        return encoder, embedding_function

    def _resolve_collection_strategy(
        self,
        *,
        client_type: str,
        partition_mode: bool,
        collection_arg: Optional[str],
    ) -> CollectionStrategy:
        """Select the collection naming strategy for the current run."""
        if partition_mode and client_type in {"http", "cloud"}:
            if collection_arg:
                return PartitionCollectionStrategy(prefix=collection_arg)
            return PartitionCollectionStrategy()
        if collection_arg is None:
            raise ValueError(
                "--collection is required for this indexing mode; supply a name."
            )
        return FixedCollectionStrategy(collection_arg)

    def _build_e2e_config(self, args: argparse.Namespace) -> E2ETestConfig:
        sample_size = max(1, int(args.e2e_sample_size))
        output_path = args.e2e_output.resolve()
        rng = (
            random.Random(int(args.e2e_random_seed))
            if args.e2e_random_seed is not None
            else random.Random()
        )
        recorder = E2ETestRecorder(output_path=output_path)
        return E2ETestConfig(
            sample_size=sample_size,
            recorder=recorder,
            rng=rng,
        )

    def _handle_single_index(self, args: argparse.Namespace) -> int:
        if args.config is None:
            logging.error("--config is required when --partition-dir is not provided")
            return 1

        client_type = str(args.client_type)
        try:
            collection_strategy = self._resolve_collection_strategy(
                client_type=client_type,
                partition_mode=False,
                collection_arg=args.collection,
            )
        except ValueError as exc:
            logging.error(str(exc))
            return 1

        target_collection = collection_strategy.collection_name(None)

        cloud_tenant: Optional[str] = None
        cloud_database: Optional[str] = None
        cloud_api_key: Optional[str] = None
        cloud_host_override = "api.trychroma.com"
        cloud_port_override = 443
        cloud_ssl = True
        e2e_config: Optional[E2ETestConfig] = None
        if args.e2e_test_run:
            if args.resume:
                logging.error("--e2e-test-run cannot be combined with --resume")
                return 1
            e2e_config = self._build_e2e_config(args)

        persist_dir: Optional[Path] = args.persist_dir
        if client_type == "persistent":
            if persist_dir is None:
                logging.error(
                    "--persist-dir is required when using the persistent Chroma client"
                )
                return 1
        else:
            if persist_dir is None:
                default_state_dir = Path.cwd() / ".vectorize_state" / target_collection
                logging.info(
                    "Using default state directory %s for resume metadata.",
                    default_state_dir,
                )
                persist_dir = default_state_dir

        api_key = args.openai_api_key or os.getenv("OPENAI_API_KEY")
        if not api_key:
            logging.error(
                "OpenAI API key is required (use --openai-api-key or set OPENAI_API_KEY)."
            )
            return 1

        try:
            encoder, embedding_function = self._initialize_embedding_resources(
                args.embedding_model, api_key
            )
        except RuntimeError as exc:
            logging.error(str(exc))
            return 1

        http_host, http_port, http_ssl, http_headers = self._resolve_http_settings(args)
        (
            cloud_tenant,
            cloud_database,
            cloud_api_key,
            cloud_host_override,
            cloud_port_override,
            cloud_ssl,
        ) = (None, None, None, "api.trychroma.com", 443, True)

        if client_type == "http" and not http_host:
            logging.error(
                "--chroma-server-host must be provided when --client-type=http"
            )
            return 1
        if client_type == "cloud":
            (
                cloud_tenant,
                cloud_database,
                cloud_api_key,
                cloud_host_override,
                cloud_port_override,
                cloud_ssl,
            ) = self._resolve_cloud_settings(args)
            if cloud_api_key is None:
                logging.error(
                    "--chroma-api-token (or CHROMA_API_TOKEN) "
                    "is required when --client-type=cloud"
                )
                return 1

        if client_type == "cloud" and args.delete_stale:
            logging.warning(
                "--delete-stale is not supported with "
                "--client-type=cloud; skipping cleanup step."
            )
        if client_type == "cloud":
            (
                cloud_tenant,
                cloud_database,
                cloud_api_key,
                cloud_host_override,
                cloud_port_override,
                cloud_ssl,
            ) = self._resolve_cloud_settings(args)
            if cloud_api_key is None:
                logging.error(
                    "--chroma-api-token (or CHROMA_API_TOKEN) "
                    "is required when --client-type=cloud"
                )
                return 1

        if persist_dir is None:
            raise RuntimeError("Persist directory could not be determined.")
        persist_dir = persist_dir.resolve()
        persist_dir.mkdir(parents=True, exist_ok=True)

        result = self._run_index_for_config(
            config_path=args.config.resolve(),
            collection_name=target_collection,
            persist_dir=persist_dir,
            batch_size=args.batch_size,
            embedding_model=args.embedding_model,
            resume=args.resume,
            skip_validation=args.skip_validation,
            embedding_function=embedding_function,
            encoder=encoder,
            resume_state_file=args.resume_state_file,
            client_type=client_type,
            http_host=http_host,
            http_port=http_port,
            http_ssl=http_ssl,
            http_headers=http_headers,
            cloud_tenant=cloud_tenant,
            cloud_database=cloud_database,
            cloud_api_key=cloud_api_key,
            cloud_host=cloud_host_override,
            cloud_port=cloud_port_override,
            cloud_ssl=cloud_ssl,
            model_registry=args.model_registry,
            truncation_strategy=args.truncation_strategy,
            token_limit=args.token_limit,
            embedding_token_limit=getattr(args, "embedding_token_limit", 8191),
            extra_metadata=None,
            model_metadata=None,
            e2e_config=e2e_config,
        )
        if result == 0 and e2e_config is not None:
            try:
                e2e_config.recorder.write()
                logging.info(
                    "E2E sample audit written to %s",
                    e2e_config.recorder.output_path,
                )
            except OSError as exc:
                logging.warning("Failed to write E2E audit file: %s", exc)
        return result

    def _handle_partition_index(self, args: argparse.Namespace) -> int:
        if args.partition_manifest and args.partition_dir:
            logging.error(
                "--partition-dir and --partition-manifest cannot be used together."
            )
            return 1
        if args.delete_stale and not args.partition_manifest:
            logging.warning(
                "--delete-stale is ignored unless --partition-manifest is provided."
            )
        if args.partition_out_dir is None:
            logging.error(
                "--partition-out-dir is required when performing partition indexing."
            )
            return 1
        if args.config is not None:
            logging.error(
                "--config cannot be used when indexing partitions "
                "(use manifest or partition directories)."
            )
            return 1
        if args.persist_dir is not None:
            logging.warning(
                "Ignoring --persist-dir because partition indexing "
                "manages persistence per partition."
            )
        if args.resume_state_file is not None:
            logging.error(
                "--resume-state-file is not supported when using partition indexing."
            )
            return 1

        partition_model_metadata: Dict[str, Dict[str, Dict[str, Any]]] = {}
        rate_limited_partitions: Deque[Tuple[str, Path, Path]] = deque()
        rate_limit_attempts: Dict[str, int] = {}
        rate_limit_backoff: Dict[str, float] = {}
        rate_limit_lock = threading.Lock()
        rate_limit_exit_code = 429
        max_rate_limit_retries = 5
        base_rate_limit_backoff = 1.0

        e2e_config: Optional[E2ETestConfig] = None
        if args.e2e_test_run:
            if args.resume:
                logging.error("--e2e-test-run cannot be combined with --resume")
                return 1
            e2e_config = self._build_e2e_config(args)

        if args.partition_manifest:
            manifest_path = args.partition_manifest
            if not manifest_path.exists():
                logging.error("Partition manifest %s does not exist", manifest_path)
                return 1
            try:
                partition_entries = load_partition_manifest_entries(
                    manifest_path.resolve(), args.partition_config_name
                )
            except (OSError, ValueError, json.JSONDecodeError) as exc:
                logging.error(
                    "Failed to load partition manifest from %s: %s",
                    manifest_path,
                    exc,
                )
                return 1
            if not partition_entries:
                logging.error(
                    "Partition manifest %s does not contain any partitions.",
                    manifest_path,
                )
                return 1
            try:
                manifest_json = json.loads(manifest_path.read_text(encoding="utf-8"))
                partitions_meta = manifest_json.get("partitions", [])
                if isinstance(partitions_meta, list):
                    for entry in partitions_meta:
                        if not isinstance(entry, Mapping):
                            continue
                        name = entry.get("name")
                        if not isinstance(name, str):
                            continue
                        models_entry = entry.get("models", {})
                        if not isinstance(models_entry, Mapping):
                            continue
                        model_meta: Dict[str, Dict[str, Any]] = {}
                        for model_name, model_info in models_entry.items():
                            if not isinstance(model_info, Mapping):
                                continue
                            schema_version = model_info.get("schema_version")
                            if isinstance(schema_version, int):
                                model_meta.setdefault(model_name, {})[
                                    "schema_version"
                                ] = schema_version
                        if model_meta:
                            partition_model_metadata[name] = model_meta
            except (OSError, json.JSONDecodeError) as exc:
                logging.warning(
                    "Failed to read manifest metadata from %s: %s",
                    manifest_path,
                    exc,
                )
            partition_origin_desc = str(manifest_path)
        else:
            if args.partition_dir is None:
                logging.error(
                    "Either --partition-dir or --partition-manifest "
                    "must be provided for partition indexing."
                )
                return 1
            if not args.partition_dir.exists() or not args.partition_dir.is_dir():
                logging.error(
                    "Partition directory %s does not exist or is not a directory",
                    args.partition_dir,
                )
                return 1
            partition_dirs = sorted(
                [entry for entry in args.partition_dir.iterdir() if entry.is_dir()],
                key=lambda path: path.name,
            )
            if not partition_dirs:
                logging.error(
                    "No partition subdirectories found in %s", args.partition_dir
                )
                return 1
            partition_entries = [
                (
                    path.name,
                    path.resolve(),
                    (path / args.partition_config_name).resolve(),
                )
                for path in partition_dirs
            ]
            partition_origin_desc = str(args.partition_dir)

        client_type = str(args.client_type)
        api_key = args.openai_api_key or os.getenv("OPENAI_API_KEY")
        if not api_key:
            logging.error(
                "OpenAI API key is required (use --openai-api-key or set OPENAI_API_KEY)."
            )
            return 1

        try:
            encoder, embedding_function = self._initialize_embedding_resources(
                args.embedding_model, api_key
            )
        except RuntimeError as exc:
            logging.error(str(exc))
            return 1

        http_host, http_port, http_ssl, http_headers = self._resolve_http_settings(args)
        (
            cloud_tenant,
            cloud_database,
            cloud_api_key,
            cloud_host_override,
            cloud_port_override,
            cloud_ssl,
        ) = (None, None, None, "api.trychroma.com", 443, True)

        if client_type == "http" and not http_host:
            logging.error(
                "--chroma-server-host must be provided when --client-type=http"
            )
            return 1
        if client_type == "cloud":
            (
                cloud_tenant,
                cloud_database,
                cloud_api_key,
                cloud_host_override,
                cloud_port_override,
                cloud_ssl,
            ) = self._resolve_cloud_settings(args)
            if cloud_api_key is None:
                logging.error(
                    "--chroma-api-token (or CHROMA_API_TOKEN) "
                    "is required when --client-type=cloud"
                )
                return 1

        if client_type == "cloud" and args.delete_stale:
            logging.warning(
                "--delete-stale is not supported with "
                "--client-type=cloud; skipping cleanup step."
            )

        try:
            collection_strategy = self._resolve_collection_strategy(
                client_type=client_type,
                partition_mode=True,
                collection_arg=args.collection,
            )
        except ValueError as exc:
            logging.error(str(exc))
            return 1

        partition_out_root = args.partition_out_dir.resolve()
        partition_out_root.mkdir(parents=True, exist_ok=True)

        parallel_partitions = int(getattr(args, "parallel_partitions", 1) or 1)
        if parallel_partitions < 1:
            logging.error("--parallel-partitions must be at least 1.")
            return 1
        if e2e_config is not None and parallel_partitions > 1:
            logging.warning(
                "Disabling parallel partition indexing because --e2e-test-run is active."
            )
            parallel_partitions = 1

        logging.info(
            "Detected %d partition(s) from %s; output will be written under %s",
            len(partition_entries),
            partition_origin_desc,
            partition_out_root,
        )

        validated_partitions: List[Tuple[str, Path, Path]] = []
        for partition_name, partition_path, config_path in partition_entries:
            if not partition_path.exists() or not partition_path.is_dir():
                logging.error(
                    "Partition %s (%s) does not exist or is not a directory",
                    partition_name,
                    partition_path,
                )
                return 1
            if not config_path.exists():
                logging.error(
                    "Partition %s is missing config file %s",
                    partition_name,
                    config_path,
                )
                return 1
            validated_partitions.append((partition_name, partition_path, config_path))

        def _index_single_partition(entry: Tuple[str, Path, Path]) -> int:
            partition_name, _partition_path, config_path = entry
            try:
                collection_name = collection_strategy.collection_name(partition_name)
            except ValueError as exc:
                logging.error(
                    "Failed to resolve collection name for partition %s: %s",
                    partition_name,
                    exc,
                )
                return 1

            partition_persist_dir = partition_out_root / partition_name
            logging.info(
                "Indexing partition %s (collection=%s, config=%s, persist_dir=%s)",
                partition_name,
                collection_name,
                config_path,
                partition_persist_dir,
            )

            if parallel_partitions == 1:
                local_encoder = encoder
                local_embedding_function = embedding_function
            else:
                try:
                    local_encoder, local_embedding_function = (
                        self._initialize_embedding_resources(
                            args.embedding_model, api_key
                        )
                    )
                except RuntimeError as exc:
                    logging.error(
                        "Failed to prepare embedding resources for partition %s: %s",
                        partition_name,
                        exc,
                    )
                    return 1

            try:
                result = self._run_index_for_config(
                    config_path=config_path,
                    collection_name=collection_name,
                    persist_dir=partition_persist_dir,
                    batch_size=args.batch_size,
                    embedding_model=args.embedding_model,
                    resume=args.resume,
                    skip_validation=args.skip_validation,
                    embedding_function=local_embedding_function,
                    encoder=local_encoder,
                    resume_state_file=None,
                    client_type=client_type,
                    http_host=http_host,
                    http_port=http_port,
                    http_ssl=http_ssl,
                    http_headers=http_headers,
                    cloud_tenant=cloud_tenant,
                    cloud_database=cloud_database,
                    cloud_api_key=cloud_api_key,
                    cloud_host=cloud_host_override,
                    cloud_port=cloud_port_override,
                    cloud_ssl=cloud_ssl,
                    model_registry=args.model_registry,
                    truncation_strategy=args.truncation_strategy,
                    token_limit=args.token_limit,
                    embedding_token_limit=getattr(args, "embedding_token_limit", 8191),
                    extra_metadata={"partition_name": partition_name},
                    model_metadata=partition_model_metadata.get(partition_name),
                    e2e_config=e2e_config,
                )
            except Exception as exc:
                if _is_rate_limit_error(exc):
                    retry_after = _extract_retry_after_seconds(exc)
                    with rate_limit_lock:
                        rate_limited_partitions.append(entry)
                        attempts = rate_limit_attempts.get(partition_name, 0) + 1
                        rate_limit_attempts[partition_name] = attempts
                        if retry_after is not None:
                            rate_limit_backoff[partition_name] = retry_after
                        else:
                            rate_limit_backoff.setdefault(
                                partition_name, base_rate_limit_backoff
                            )
                    if retry_after is not None:
                        logging.warning(
                            "Partition %s hit OpenAI rate limit "
                            "(attempt %d, retry in %.2fs); deferring.",
                            partition_name,
                            attempts,
                            retry_after,
                        )
                    else:
                        logging.warning(
                            "Partition %s hit OpenAI rate limit "
                            "(attempt %d); deferring for later retry.",
                            partition_name,
                            attempts,
                        )
                    return rate_limit_exit_code
                raise

            if result == 0:
                with rate_limit_lock:
                    rate_limit_attempts.pop(partition_name, None)
                    rate_limit_backoff.pop(partition_name, None)
            return result

        def _retry_rate_limited_partitions() -> bool:
            attempt_cycle = 1
            while True:
                with rate_limit_lock:
                    pending_entries = list(rate_limited_partitions)
                    pending_backoffs = {
                        entry[0]: rate_limit_backoff.get(entry[0], 0.0)
                        for entry in pending_entries
                    }
                    rate_limited_partitions.clear()
                    for name in pending_backoffs:
                        rate_limit_backoff.pop(name, None)
                if not pending_entries:
                    return True

                max_delay = max(pending_backoffs.values(), default=0.0)
                if max_delay <= 0.0:
                    max_delay = min(
                        30.0, base_rate_limit_backoff * max(1, attempt_cycle - 1)
                    )
                else:
                    max_delay = min(30.0, max_delay)
                if max_delay > 0:
                    logging.info(
                        "Sleeping %.2f second(s) before retrying %d "
                        "rate-limited partition(s) (cycle %d).",
                        max_delay,
                        len(pending_entries),
                        attempt_cycle,
                    )
                    time.sleep(max_delay)

                logging.info(
                    "Retrying %d partition(s) deferred by rate limiting (cycle %d).",
                    len(pending_entries),
                    attempt_cycle,
                )

                progress = False
                for entry in pending_entries:
                    partition_name = entry[0]
                    attempts = rate_limit_attempts.get(partition_name, 0)
                    if attempts > max_rate_limit_retries:
                        logging.error(
                            "Partition %s exceeded maximum rate limit retries (%d).",
                            partition_name,
                            max_rate_limit_retries,
                        )
                        return False
                    result = _index_single_partition(entry)
                    if result == rate_limit_exit_code:
                        continue
                    if result != 0:
                        logging.error(
                            "Partition %s failed with exit code %s during rate limit retry.",
                            partition_name,
                            result,
                        )
                        return False
                    progress = True
                attempt_cycle += 1
                if progress:
                    continue

        if parallel_partitions == 1:
            for entry in validated_partitions:
                result = _index_single_partition(entry)
                if result == rate_limit_exit_code:
                    continue
                if result != 0:
                    return result
        else:
            logging.info(
                "Parallel partition indexing enabled with %d worker(s).",
                parallel_partitions,
            )
            had_error = False
            with ThreadPoolExecutor(max_workers=parallel_partitions) as executor:
                future_map = {
                    executor.submit(_index_single_partition, entry): entry[0]
                    for entry in validated_partitions
                }
                for future in as_completed(future_map):
                    partition_name = future_map[future]
                    try:
                        result = future.result()
                    except Exception:  # pragma: no cover - defensive
                        logging.exception(
                            "Partition %s raised an unexpected exception",
                            partition_name,
                        )
                        had_error = True
                        continue
                    if result == rate_limit_exit_code:
                        logging.info(
                            "Partition %s deferred due to rate limiting; will retry later.",
                            partition_name,
                        )
                        continue
                    if result != 0:
                        logging.error(
                            "Partition %s failed with exit code %s",
                            partition_name,
                            result,
                        )
                        had_error = True
            if had_error:
                return 1

        if not _retry_rate_limited_partitions():
            return 1

        logging.info(
            "Completed indexing for %d partition(s) from %s",
            len(partition_entries),
            partition_origin_desc,
        )

        if e2e_config is not None:
            try:
                e2e_config.recorder.write()
                logging.info(
                    "E2E sample audit written to %s",
                    e2e_config.recorder.output_path,
                )
            except OSError as exc:
                logging.warning("Failed to write E2E audit file: %s", exc)

        if args.partition_manifest and args.delete_stale:
            self._delete_stale_partitions(
                manifest_path=args.partition_manifest.resolve(),
                partition_out_root=partition_out_root,
                client_type=client_type,
                collection_strategy=collection_strategy,
                http_host=http_host,
                http_port=http_port,
                http_ssl=http_ssl,
                http_headers=http_headers,
            )
        return 0

    def _delete_stale_partitions(
        self,
        *,
        manifest_path: Path,
        partition_out_root: Path,
        client_type: str,
        collection_strategy: CollectionStrategy,
        http_host: Optional[str],
        http_port: Optional[int],
        http_ssl: bool,
        http_headers: Mapping[str, str],
    ) -> None:
        try:
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
            if not isinstance(manifest, dict):
                raise ValueError("Manifest must contain a JSON object")
        except (OSError, ValueError, json.JSONDecodeError) as exc:
            logging.error(
                "Unable to load manifest %s for stale deletion: %s", manifest_path, exc
            )
            return

        partitions = manifest.get("partitions", [])
        if not isinstance(partitions, list):
            logging.warning(
                "Manifest partitions entry is invalid; skipping stale deletion."
            )
            return

        stale_entries: List[Dict[str, Any]] = [
            entry
            for entry in partitions
            if isinstance(entry, dict)
            and entry.get("stale")
            and not entry.get("deleted")
        ]
        if not stale_entries:
            logging.info("No stale partitions marked for deletion.")
            return

        timestamp = (
            datetime.now(timezone.utc)
            .replace(tzinfo=None)
            .isoformat(timespec="seconds")
        )
        deleted_names: List[str] = []

        if client_type == "persistent":
            for entry in stale_entries:
                name = entry.get("name")
                if not name:
                    continue
                target_dir = partition_out_root / name
                if target_dir.exists():
                    try:
                        shutil.rmtree(target_dir)
                        logging.info("Removed stale partition store at %s", target_dir)
                    except OSError as exc:
                        logging.warning(
                            "Failed to remove stale partition directory %s (%s)",
                            target_dir,
                            exc,
                        )
                        continue
                entry["deleted"] = True
                entry["deleted_at"] = timestamp
                deleted_names.append(name)
        else:
            try:
                client = self._build_chroma_client(
                    client_type=client_type,
                    persist_dir=None,
                    http_host=http_host,
                    http_port=http_port,
                    http_ssl=http_ssl,
                    http_headers=http_headers,
                )
            except Exception as exc:  # pragma: no cover - defensive
                logging.error(
                    "Failed to prepare Chroma client for stale deletion: %s", exc
                )
                return
            if isinstance(collection_strategy, PartitionCollectionStrategy):
                for entry in stale_entries:
                    name = entry.get("name")
                    if not name:
                        continue
                    try:
                        target_collection = collection_strategy.collection_name(name)
                    except ValueError:
                        logging.warning(
                            "Unable to resolve collection name "
                            "for stale partition %s; skipping.",
                            name,
                        )
                        continue
                    try:
                        client.delete_collection(name=target_collection)
                        logging.info(
                            "Deleted collection %s for stale partition '%s'",
                            target_collection,
                            name,
                        )
                    except Exception as exc:  # pragma: no cover - defensive
                        logging.warning(
                            "Failed to delete collection %s for partition %s (%s)",
                            target_collection,
                            name,
                            exc,
                        )
                        continue
                    entry["deleted"] = True
                    entry["deleted_at"] = timestamp
                    deleted_names.append(name)
            else:
                collection_name = collection_strategy.collection_name(None)
                try:
                    collection = client.get_collection(name=collection_name)
                except Exception as exc:  # pragma: no cover - defensive
                    logging.error(
                        "Failed to open collection %s for stale deletion: %s",
                        collection_name,
                        exc,
                    )
                    return

                for entry in stale_entries:
                    name = entry.get("name")
                    if not name:
                        continue
                    try:
                        collection.delete(where={"partition_name": name})
                        logging.info(
                            "Deleted stale partition '%s' from collection %s",
                            name,
                            collection_name,
                        )
                    except Exception as exc:  # pragma: no cover - defensive
                        logging.warning(
                            "Failed to delete stale partition '%s' from collection %s (%s)",
                            name,
                            collection_name,
                            exc,
                        )
                        continue
                    entry["deleted"] = True
                    entry["deleted_at"] = timestamp
                    deleted_names.append(name)

        if not deleted_names:
            logging.info("No stale partitions were deleted.")
            return

        try:
            manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
        except OSError as exc:
            logging.warning(
                "Failed to update manifest %s after deletion: %s", manifest_path, exc
            )
        else:
            logging.info(
                "Deleted %d stale partition(s): %s",
                len(deleted_names),
                ", ".join(sorted(deleted_names)),
            )

    def handle_index(self, args: argparse.Namespace) -> int:
        if args.partition_dir or args.partition_manifest:
            return self._handle_partition_index(args)
        return self._handle_single_index(args)

    def handle_drop_models(self, args: argparse.Namespace) -> int:
        logging.basicConfig(
            level=getattr(logging, str(args.log_level).upper(), logging.INFO),
            format="%(levelname)s %(message)s",
        )

        model_registry = getattr(args, "model_registry", None)
        if model_registry is None:
            logging.error(
                "Model registry not provided; ensure --model is specified when "
                "invoking the CLI."
            )
            return 1

        try:
            from idxr.prepare_datasets_lib.config import (
                load_drop_config,
            )
        except ImportError as exc:  # pragma: no cover - defensive
            logging.error("Unable to load drop configuration utilities: %s", exc)
            return 1

        config_path = args.config.resolve()
        try:
            drop_configs, _ = load_drop_config(
                config_path, model_registry=model_registry
            )
        except (OSError, ValueError, KeyError, json.JSONDecodeError) as exc:
            logging.error("Failed to load drop configuration: %s", exc)
            return 1

        if not drop_configs:
            logging.warning("Drop configuration does not reference any models.")
            return 0

        filters: List[Tuple[str, Dict[str, Any], List[str]]] = []
        for model_name, cfg in drop_configs.items():
            partitions = list(getattr(cfg, "partitions", []) or [])
            schema_versions = list(getattr(cfg, "schema_versions", []) or [])
            clauses: List[Dict[str, Any]] = [{"model_name": model_name}]
            if partitions:
                clauses.append({"partition_name": {"$in": partitions}})
            if schema_versions:
                clauses.append({"schema_version": {"$in": schema_versions}})
            if len(clauses) == 1:
                where = clauses[0]
            else:
                where = {"$and": clauses}
            filters.append((model_name, where, partitions))

        dry_run = not args.apply
        if dry_run:
            logging.info("Dry run – no data will be deleted. Use --apply to execute.")
        for model_name, where, partitions in filters:
            logging.info(
                "Model %s -> where=%s%s",
                model_name,
                where,
                (f" (partitions {partitions})" if partitions else ""),
            )

        if dry_run:
            return 0

        client_type = str(args.client_type)
        http_host, http_port, http_ssl, http_headers = self._resolve_http_settings(args)
        persist_dir: Optional[Path] = args.persist_dir

        if client_type == "persistent":
            if persist_dir is None:
                logging.error(
                    "--persist-dir is required when using the persistent Chroma client."
                )
                return 1
            persist_dir = persist_dir.resolve()
        else:
            if not http_host:
                logging.error(
                    "--chroma-server-host must be provided when --client-type=http"
                )
                return 1

        try:
            client = self._build_chroma_client(
                client_type=client_type,
                persist_dir=persist_dir,
                http_host=http_host,
                http_port=http_port,
                http_ssl=http_ssl,
                http_headers=http_headers,
            )
            collection = client.get_collection(name=args.collection)
        except Exception as exc:  # pragma: no cover - defensive
            logging.error("Failed to open collection %s: %s", args.collection, exc)
            return 1

        for model_name, where, partitions in filters:
            try:
                result = collection.get(where=where, include=[]) or {}
            except AttributeError:
                try:
                    result = collection.get(where=where) or {}
                except Exception as exc:  # pragma: no cover - defensive
                    logging.error(
                        "Failed to fetch documents for %s (filter=%s): %s",
                        model_name,
                        where,
                        exc,
                    )
                    return 1
            except Exception as exc:  # pragma: no cover - defensive
                logging.error(
                    "Failed to fetch documents for %s (filter=%s): %s",
                    model_name,
                    where,
                    exc,
                )
                return 1

            ids = list(result.get("ids") or [])  # type: ignore[arg-type]
            if not ids:
                logging.info(
                    "No documents matched drop filter for %s (partitions=%s)",
                    model_name,
                    partitions or "all",
                )
                continue

            try:
                collection.delete(ids=ids)
            except Exception as exc:  # pragma: no cover - defensive
                logging.error(
                    "Failed to delete documents for %s (ids=%s): %s",
                    model_name,
                    ids,
                    exc,
                )
                return 1
            logging.info(
                "Deleted %s document(s) for %s (%s partition filter).",
                format_int(len(ids)),
                model_name,
                format_int(len(partitions)) if partitions else "all",
            )

        if args.partition_manifest:
            try:
                from idxr.prepare_datasets_lib.drop import (
                    apply_drop_manifest,
                )
            except ImportError as exc:  # pragma: no cover - defensive
                logging.warning(
                    "Could not update manifest %s (%s)", args.partition_manifest, exc
                )
            else:
                apply_drop_manifest(
                    manifest_path=args.partition_manifest.resolve(),
                    drop_config_path=config_path,
                    apply_changes=True,
                    remove_local=False,
                    performed_by=args.performed_by,
                    model_registry=model_registry,
                )
        return 0

    def handle_display(self, args: argparse.Namespace) -> int:
        if args.chunk_size <= 0:
            logging.error("chunk-size must be a positive integer")
            return 1

        persist_dir: Path = args.persist_dir
        if not persist_dir.exists():
            logging.warning(
                "Persist directory %s does not exist yet; "
                "attempting to open collection anyway.",
                persist_dir,
            )

        client = chromadb.PersistentClient(path=str(persist_dir))
        try:
            collection = client.get_collection(name=args.collection)
        except (InvalidCollectionError, ValueError) as exc:
            logging.error(
                "Collection '%s' could not be opened: %s", args.collection, exc
            )
            return 1

        total, counts, missing = summarize_collection(
            collection,
            chunk_size=args.chunk_size,
            log_progress=True,
        )
        print(f"Collection: {args.collection}")
        print(f"Persist dir: {persist_dir}")
        print(f"Total documents: {format_int(total)}")

        if total == 0:
            print("No documents indexed yet.")
            return 0

        if not counts and missing == 0:
            print("No metadata available to summarize.")
            return 0

        rows = sorted(counts.items(), key=lambda item: item[1], reverse=True)
        distinct_models = len(rows)
        print(f"Distinct models: {distinct_models}")
        top = args.top
        if top == 0 or top >= len(rows):
            top_rows = rows
        else:
            top_rows = rows[:top]
        if top_rows:
            print("")
            print("Per-model document counts:")
            for model_name, count in top_rows:
                print(f"  {model_name}: {format_int(count)}")
        if top and len(rows) > top:
            remaining_count = sum(count for _, count in rows[top:])
            remaining_models = len(rows) - top
            print(
                f"  ... {remaining_models} additional model(s) totaling "
                f"{format_int(remaining_count)} document(s)"
            )
        if missing:
            print(f"Documents missing model metadata: {format_int(missing)}")
        return 0

    def handle_init_config(self, args: argparse.Namespace) -> int:
        output_path: Path = args.output
        if output_path.exists() and not args.force:
            logging.error(
                "Refusing to overwrite existing file %s (use --force)", output_path
            )
            return 1
        stub = generate_stub_config(args.model_registry)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(json.dumps(stub, indent=2), encoding="utf-8")
        logging.info("Wrote stub configuration to %s", output_path)
        return 0

    def handle_status(self, args: argparse.Namespace) -> int:
        """Display current indexing progress by analyzing resume states and errors."""
        partition_out_dir: Path = args.partition_out_dir
        partition_dir: Optional[Path] = args.partition_dir
        partition_config_name: str = args.partition_config_name

        if not partition_out_dir.exists():
            logging.error(
                "Partition output directory %s does not exist", partition_out_dir
            )
            return 1

        # Find all partition directories in output
        partition_dirs = sorted(
            [d for d in partition_out_dir.iterdir() if d.is_dir()], key=lambda p: p.name
        )

        if not partition_dirs:
            print(f"No partitions found in {partition_out_dir}")
            return 0

        print(f"\n{'='*80}")
        print("Indexing Status Report")
        print(f"{'='*80}")
        print(f"Partition Directory: {partition_out_dir}")
        if partition_dir:
            print(f"Config Directory: {partition_dir}")
        print(f"Total Partitions: {len(partition_dirs)}")
        print(f"{'='*80}\n")

        # Collect status for each partition
        partition_statuses = []
        for partition_out in partition_dirs:
            partition_name = partition_out.name

            # Try to find corresponding config directory
            config_path: Optional[Path] = None
            if partition_dir and partition_dir.exists():
                potential_config_dir = partition_dir / partition_name
                if potential_config_dir.exists():
                    config_file = potential_config_dir / partition_config_name
                    if config_file.exists():
                        config_path = config_file

            status_info = self._analyze_partition_status(
                partition_out, partition_name, config_path
            )
            partition_statuses.append(status_info)

        # Display summary
        finished_count = sum(1 for s in partition_statuses if s["status"] == "finished")
        started_count = sum(1 for s in partition_statuses if s["status"] == "started")
        errored_count = sum(1 for s in partition_statuses if s["status"] == "errored")
        not_started_count = sum(
            1 for s in partition_statuses if s["status"] == "not_started"
        )

        print("Summary:")
        print(f"  Finished:    {finished_count:>4} partition(s)")
        print(f"  Started:     {started_count:>4} partition(s)")
        print(f"  Errored:     {errored_count:>4} partition(s)")
        print(f"  Not Started: {not_started_count:>4} partition(s)")
        print(f"\n{'='*80}\n")

        # Display detailed status for each partition
        for status_info in partition_statuses:
            self._display_partition_status(status_info)

        return 0

    def _analyze_partition_status(
        self,
        partition_dir: Path,
        partition_name: str,
        config_path: Optional[Path] = None,
    ) -> Dict[str, Any]:
        """Analyze a single partition's status."""
        status_info: Dict[str, Any] = {
            "partition_name": partition_name,
            "partition_dir": partition_dir,
            "status": "not_started",
            "models": {},
            "configured_models": set(),
            "errors": [],
            "total_collection_count": 0,
        }

        # Load configured models from vectorize_config.json if available
        if config_path and config_path.exists():
            try:
                config_data = json.loads(config_path.read_text(encoding="utf-8"))
                if isinstance(config_data, dict):
                    for model_name, model_config in config_data.items():
                        if isinstance(model_config, dict) and model_config.get("path"):
                            status_info["configured_models"].add(model_name)
            except (OSError, json.JSONDecodeError) as exc:
                logging.warning("Failed to read config from %s: %s", config_path, exc)

        # Find resume state files
        resume_files = list(partition_dir.glob("*_resume_state.json"))

        if resume_files:
            # Load resume state
            for resume_file in resume_files:
                try:
                    resume_data = json.loads(resume_file.read_text(encoding="utf-8"))
                    if isinstance(resume_data, dict):
                        for model_name, model_state in resume_data.items():
                            if isinstance(model_state, dict):
                                collection_total = model_state.get(
                                    "collection_count",
                                    model_state.get("documents_indexed", 0),
                                )
                                status_info["models"][model_name] = {
                                    "complete": model_state.get("complete", False),
                                    "documents_indexed": model_state.get(
                                        "documents_indexed", 0
                                    ),
                                    "collection_count": collection_total,
                                    "indexed_at": model_state.get("indexed_at", ""),
                                    "row_index": model_state.get("row_index", 0),
                                    "started": model_state.get("started", False),
                                }
                except (OSError, json.JSONDecodeError) as exc:
                    logging.warning(
                        "Failed to read resume state from %s: %s", resume_file, exc
                    )
            status_info["total_collection_count"] = sum(
                model_details.get("collection_count", 0)
                for model_details in status_info["models"].values()
            )

        # Check for errors
        errors_dir = partition_dir / "errors"
        max_error_row_by_model: Dict[str, int] = {}
        max_error_row_any = -1
        if errors_dir.exists() and errors_dir.is_dir():
            error_files = sorted(errors_dir.glob("*.yaml"), reverse=True)
            for index, error_file in enumerate(error_files):
                try:
                    import yaml

                    error_data = yaml.safe_load(error_file.read_text(encoding="utf-8"))
                    if isinstance(error_data, dict):
                        if index < 5:  # Show up to 5 most recent errors
                            status_info["errors"].append(
                                {
                                    "timestamp": error_data.get("timestamp", ""),
                                    "model": error_data.get("model", ""),
                                    "exception_type": error_data.get(
                                        "exception", {}
                                    ).get("type", ""),
                                    "exception_message": error_data.get(
                                        "exception", {}
                                    ).get("message", "")[:100],
                                    "file": error_file.name,
                                }
                            )

                        candidate_rows: List[int] = []
                        rows_block = error_data.get("rows")
                        if isinstance(rows_block, list):
                            for row_entry in rows_block:
                                if isinstance(row_entry, Mapping):
                                    row_idx = row_entry.get("row_index")
                                    coerced_row = _coerce_optional_int(row_idx)
                                    if coerced_row is not None:
                                        candidate_rows.append(coerced_row)
                        resume_block = error_data.get("resume_state")
                        if isinstance(resume_block, Mapping):
                            resume_row = resume_block.get("row_index")
                            coerced_resume_row = _coerce_optional_int(resume_row)
                            if coerced_resume_row is not None:
                                candidate_rows.append(coerced_resume_row)

                        if candidate_rows:
                            max_row = max(candidate_rows)
                            model_name = error_data.get("model")
                            if isinstance(model_name, str) and model_name:
                                previous_max = max_error_row_by_model.get(model_name)
                                if previous_max is None or max_row > previous_max:
                                    max_error_row_by_model[model_name] = max_row
                            if max_row > max_error_row_any:
                                max_error_row_any = max_row
                except Exception as exc:
                    logging.warning("Failed to read error file %s: %s", error_file, exc)

        # Determine overall status
        model_names = set(status_info["models"].keys())
        all_models_complete = bool(status_info["models"]) and all(
            model_state.get("complete", False)
            for model_state in status_info["models"].values()
        )
        if status_info["configured_models"]:
            all_models_complete = all_models_complete and status_info[
                "configured_models"
            ].issubset(model_names)

        if all_models_complete:
            status_info["status"] = "finished"
            return status_info

        progressed_past_errors = False
        if status_info["errors"]:
            model_progress_rows: Dict[str, int] = {}
            for model_name, model_state in status_info["models"].items():
                row_index_value = model_state.get("row_index")
                coerced_progress_row = _coerce_optional_int(row_index_value)
                if coerced_progress_row is not None:
                    model_progress_rows[model_name] = coerced_progress_row

            if max_error_row_by_model:
                progressed_past_errors = True
                for model_name, max_error_row in max_error_row_by_model.items():
                    progress_row = model_progress_rows.get(model_name)
                    if progress_row is None or progress_row <= max_error_row:
                        progressed_past_errors = False
                        break
            elif max_error_row_any >= 0 and model_progress_rows:
                progressed_past_errors = (
                    max(model_progress_rows.values(), default=-1) > max_error_row_any
                )

            if progressed_past_errors:
                status_info["status"] = "started"
            else:
                status_info["status"] = "errored"
            return status_info

        if status_info["models"]:
            configured_models = status_info["configured_models"]
            if configured_models:
                any_started = any(
                    status_info["models"].get(model, {}).get("started", False)
                    for model in configured_models
                )
                status_info["status"] = "started" if any_started else "not_started"
            else:
                any_started = any(
                    model_state.get("started", False)
                    for model_state in status_info["models"].values()
                )
                status_info["status"] = "started" if any_started else "not_started"
        else:
            status_info["status"] = "not_started"

        return status_info

    def _display_partition_status(self, status_info: Dict[str, Any]) -> None:
        """Display status information for a single partition."""
        partition_name = status_info["partition_name"]
        status = status_info["status"]

        # Status indicator
        status_emoji = {
            "finished": "✓",
            "started": "→",
            "errored": "✗",
            "not_started": "○",
        }

        status_color = {
            "finished": "FINISHED",
            "started": "STARTED",
            "errored": "ERRORED",
            "not_started": "NOT STARTED",
        }

        indicator = status_emoji.get(status, "?")
        status_text = status_color.get(status, status.upper())

        print(f"[{indicator}] {partition_name:<25} {status_text}")

        # Show configured models count if available
        configured_models = status_info.get("configured_models", set())
        if configured_models:
            models_indexed = len(status_info["models"])
            models_configured = len(configured_models)
            print(f"    Models: {models_indexed}/{models_configured} started")
        elif status_info["models"]:
            print(f"    Models: {len(status_info['models'])}")

        total_collection_count = status_info.get("total_collection_count", 0)
        if total_collection_count > 0:
            print(
                f"    Total Documents (collection count): {format_int(total_collection_count)}"
            )

        # Show model details if available
        if status_info["models"]:
            for model_name, model_state in sorted(status_info["models"].items()):
                complete_marker = "✓" if model_state.get("complete") else "→"
                total_docs = model_state.get(
                    "collection_count",
                    model_state.get("documents_indexed", 0),
                )
                resume_run_docs = model_state.get("documents_indexed", 0)
                indexed_at = model_state.get("indexed_at", "")
                print(
                    f"      [{complete_marker}] {model_name}: "
                    f"{format_int(total_docs)} docs total"
                )
                if resume_run_docs and resume_run_docs != total_docs:
                    print(
                        f"          Last run indexed: {format_int(resume_run_docs)} docs"
                    )
                if indexed_at:
                    print(f"          Last indexed: {indexed_at}")

        # Show configured but not started models
        if configured_models:
            not_started_models = configured_models - set(status_info["models"].keys())
            if not_started_models:
                print(f"    Not Started: {len(not_started_models)} model(s)")
                for model_name in sorted(not_started_models):
                    print(f"      [○] {model_name}: not started")

        # Show errors if present
        if status_info["errors"]:
            print(f"    Errors: {len(status_info['errors'])} error report(s)")
            for error in status_info["errors"][:3]:  # Show up to 3 most recent
                print(f"      • {error['timestamp']} - {error['model']}")
                print(
                    f"        {error['exception_type']}: {error['exception_message']}"
                )

        print()

    def handle_generate_query_config(self, args: argparse.Namespace) -> int:
        """Generate query configuration from partition resume states."""
        partition_out_dir: Path = args.partition_out_dir
        output_path: Path = args.output
        collection_prefix: Optional[str] = args.collection_prefix

        # Environment generation args
        generate_environment: bool = args.generate_environment
        environment_output_path: Optional[Path] = args.environment_output

        if not partition_out_dir.exists():
            logging.error(
                "Partition output directory %s does not exist", partition_out_dir
            )
            return 1

        # Validate environment args if requested
        if generate_environment:
            if not args.model_registry_target:
                logging.error(
                    "--model-registry-target is required when --generate-environment is used"
                )
                return 1

            if args.chroma_client_type == "cloud":
                if not (args.cloud_tenant and args.cloud_database):
                    logging.error(
                        "--cloud-tenant and --cloud-database are required for cloud client type"
                    )
                    return 1
            elif args.chroma_client_type == "http":
                if not (args.http_host and args.http_port):
                    logging.error(
                        "--http-host and --http-port are required for http client type"
                    )
                    return 1

        logging.info("Generating query config from %s", partition_out_dir)

        try:
            query_config = generate_query_config(
                partition_out_dir=partition_out_dir,
                output_path=output_path,
                collection_prefix=collection_prefix,
                generate_environment=generate_environment,
                environment_output_path=environment_output_path,
                embedding_model=args.embedding_model,
                model_registry_target=args.model_registry_target,
                discriminator_field=args.discriminator_field,
                chroma_client_type=args.chroma_client_type,
                cloud_api_key=args.cloud_api_key,
                cloud_tenant=args.cloud_tenant,
                cloud_database=args.cloud_database,
                http_host=args.http_host,
                http_port=args.http_port,
                http_ssl=args.http_ssl,
            )

            # Display summary
            metadata = query_config["metadata"]
            model_to_collections = query_config["model_to_collections"]

            print(f"\n{'='*80}")
            print("Query Configuration Generated")
            print(f"{'='*80}")
            print(f"Output File:        {output_path}")
            print(f"Total Models:       {metadata['total_models']}")
            print(f"Total Collections:  {metadata['total_collections']}")
            print(f"Generated At:       {metadata['generated_at']}")
            print(f"{'='*80}\n")

            # Show model -> collections mapping
            if model_to_collections:
                print("Model to Collections Mapping:")
                for model_name in sorted(model_to_collections.keys()):
                    info = model_to_collections[model_name]
                    print(f"\n  {model_name}:")
                    print(f"    Collections:     {len(info['collections'])}")
                    print(f"    Total Documents: {format_int(info['total_documents'])}")
                    print(f"    Partitions:      {len(info['partitions'])}")
                    if len(info["collections"]) <= 5:
                        for coll in info["collections"]:
                            print(f"      - {coll}")
                    else:
                        for coll in info["collections"][:3]:
                            print(f"      - {coll}")
                        print(f"      ... and {len(info['collections']) - 3} more")

            print(f"\n{'='*80}")
            logging.info("Query config successfully written to %s", output_path)

            # Display environment summary if generated
            if generate_environment:
                env_path = environment_output_path or (
                    partition_out_dir / "environment.json"
                )
                print(f"\n{'='*80}")
                print("Environment Configuration Generated")
                print(f"{'='*80}")
                print(f"Output File:        {env_path}")
                print(f"Embedding Model:    {args.embedding_model}")
                print(f"Model Registry:     {args.model_registry_target}")
                print(f"Discriminator:      {args.discriminator_field}")
                print(f"Client Type:        {args.chroma_client_type}")
                if args.chroma_client_type == "cloud":
                    print(f"Cloud Tenant:       {args.cloud_tenant}")
                    print(f"Cloud Database:     {args.cloud_database}")
                    print(
                        f"Cloud API Key:      {args.cloud_api_key or 'env:CHROMA_API_KEY'}"
                    )
                elif args.chroma_client_type == "http":
                    print(f"HTTP Host:          {args.http_host}")
                    print(f"HTTP Port:          {args.http_port}")
                    print(f"HTTP SSL:           {args.http_ssl}")
                elif args.chroma_client_type == "persistent":
                    print(f"Persist Directory:  {partition_out_dir}")
                print(f"{'='*80}\n")
                logging.info("Environment config successfully written to %s", env_path)

            return 0

        except (ValueError, OSError) as exc:
            logging.error("Failed to generate query config: %s", exc)
            return 1


def main(argv: Optional[Sequence[str]] = None) -> int:
    cli = VectorizeCLI()
    return cli.run(list(argv or []))
