# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
In-process runmode.

Intended to run locally in a plain Python process with minimal assumptions
about surrounding framework/runtime.
"""

from __future__ import annotations

import glob
import json
import multiprocessing
import os
import re
import time
import uuid
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from datetime import datetime, timezone
from io import BytesIO
from collections.abc import Callable, Iterator
from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union

from nemo_retriever.params import CaptionParams


import pandas as pd
from nemo_retriever.chart.chart_detection import graphic_elements_ocr_page_elements
from nemo_retriever.dedup.dedup import dedup_images
from nemo_retriever.model.local import NemotronOCRV1, NemotronPageElementsV3, NemotronParseV12
from nemo_retriever.ocr.ocr import _crop_b64_image_by_norm_bbox, nemotron_parse_page_elements, ocr_page_elements
from nemo_retriever.page_elements import detect_page_elements_v3
from nemo_retriever.params import DedupParams
from nemo_retriever.table.table_detection import table_structure_ocr_page_elements
from nemo_retriever.text_embed.main_text_embed import TextEmbeddingConfig, create_text_embeddings_for_df

try:
    from tqdm.auto import tqdm
except Exception:  # pragma: no cover
    tqdm = None  # type: ignore[assignment]

try:
    import pypdfium2 as pdfium
except Exception as e:  # pragma: no cover
    pdfium = None  # type: ignore[assignment]
    _PDFIUM_IMPORT_ERROR = e

from ..image.load import SUPPORTED_IMAGE_EXTENSIONS
from ..utils.convert import SUPPORTED_EXTENSIONS, convert_to_pdf_bytes
from ..ingestor import Ingestor
from ..params import ASRParams
from ..params import AudioChunkParams
from ..params import EmbedParams
from ..params.models import IMAGE_MODALITIES
from ..params import ExtractParams
from ..params import HtmlChunkParams
from ..params import IngestExecuteParams
from ..params import TextChunkParams
from ..params import VdbUploadParams
from ..pdf.extract import pdf_extraction
from ..pdf.split import _split_pdf_to_single_page_bytes, pdf_path_to_pages_df
from ..utils.remote_auth import resolve_remote_api_key
from ..txt import txt_file_to_chunks_df
from ..html import html_file_to_chunks_df

_CONTENT_COLUMNS = ("table", "chart", "infographic")


from nemo_retriever.params.utils import coerce_params as _coerce_params


def _combine_text_with_content(row, text_column, content_columns):
    """Combine page text with OCR content text for embedding."""
    parts = []
    base = row.get(text_column)
    if isinstance(base, str) and base.strip():
        parts.append(base.strip())
    for col in content_columns:
        content_list = row.get(col)
        if isinstance(content_list, list):
            for item in content_list:
                if isinstance(item, dict):
                    t = item.get("text", "")
                    if isinstance(t, str) and t.strip():
                        parts.append(t.strip())
    return "\n\n".join(parts) if parts else ""


def _deep_copy_row(row_dict: Dict[str, Any]) -> Dict[str, Any]:
    """Shallow-copy a row dict but deep-copy any nested mutable values.

    This avoids the problem where multiple exploded rows share the same
    metadata dict reference, causing later embedding writes to overwrite
    earlier ones.  We only deep-copy dicts and lists (the common mutable
    types in row data); scalars and strings are immutable and safe to share.
    """
    import copy

    out: Dict[str, Any] = {}
    for k, v in row_dict.items():
        if isinstance(v, (dict, list)):
            out[k] = copy.deepcopy(v)
        else:
            out[k] = v
    return out


def explode_content_to_rows(
    batch_df: Any,
    *,
    text_column: str = "text",
    content_columns: Sequence[str] = _CONTENT_COLUMNS,
    modality: str = "text",
    text_elements_modality: Optional[str] = None,
    structured_elements_modality: Optional[str] = None,
) -> Any:
    """Expand each page row into multiple rows for per-element embedding.

    For each row in *batch_df*:
    - One row with the original page text (``text_column``).
    - One additional row per table / chart / infographic item whose
      ``"text"`` field is non-empty.  The item text replaces
      ``text_column`` so the downstream embedding stage embeds each
      element independently.

    This mirrors the nv-ingest pipeline where every structural element
    gets its own embedding vector, improving recall for queries that
    target specific tables or charts.

    If a row has no page text *and* no structured content, the original
    row is preserved as-is to maintain batch shape.

    When a row's effective modality is ``"image"`` or ``"text_image"``,
    the page image (or a cropped region for structured content) is
    carried in a new ``_image_b64`` column so that the downstream
    embedding stage can use it for multimodal embeddings.

    Each exploded row is tagged with ``_embed_modality`` so that
    ``embed_text_main_text_embed`` can embed different element types
    with different methods (e.g. text-only for page text but
    text+image for tables).

    Parameters
    ----------
    modality : str
        Base modality for all element types (default ``"text"``).
    text_elements_modality : str | None
        Override modality for page-text rows.  Falls back to *modality*.
    structured_elements_modality : str | None
        Override modality for table/chart/infographic rows.  Falls back
        to *modality*.
    """
    text_mod = text_elements_modality or modality
    struct_mod = structured_elements_modality or modality

    if not isinstance(batch_df, pd.DataFrame) or batch_df.empty:
        return batch_df

    # Whether *any* element type requires images.
    any_images = text_mod in IMAGE_MODALITIES or struct_mod in IMAGE_MODALITIES

    # Fast path: if none of the content columns exist there is nothing to explode.
    if not any(c in batch_df.columns for c in content_columns):
        batch_df = batch_df.copy()
        if text_mod in IMAGE_MODALITIES and "page_image" in batch_df.columns:
            batch_df["_image_b64"] = batch_df["page_image"].apply(
                lambda pi: pi.get("image_b64") if isinstance(pi, dict) else None
            )
        batch_df["_embed_modality"] = text_mod
        return batch_df

    new_rows: List[Dict[str, Any]] = []
    for _, row in batch_df.iterrows():
        row_dict = row.to_dict()
        exploded_any = False

        # Extract page-level image b64 once per source row.
        page_image = row_dict.get("page_image")
        page_image_b64: Optional[str] = None
        if any_images and isinstance(page_image, dict):
            page_image_b64 = page_image.get("image_b64")

        # Row for page text.
        page_text = row_dict.get(text_column)
        if isinstance(page_text, str) and page_text.strip():
            page_row = _deep_copy_row(row_dict)
            page_row["_embed_modality"] = text_mod
            page_row["_content_type"] = "text"
            if text_mod in IMAGE_MODALITIES:
                page_row["_image_b64"] = page_image_b64
            new_rows.append(page_row)
            exploded_any = True

        # One row per structured content item.
        for col in content_columns:
            content_list = row_dict.get(col)
            if not isinstance(content_list, list):
                continue
            for item in content_list:
                if not isinstance(item, dict):
                    continue
                t = item.get("text", "")
                if not isinstance(t, str) or not t.strip():
                    continue
                content_row = _deep_copy_row(row_dict)
                content_row[text_column] = t.strip()
                content_row["_embed_modality"] = struct_mod
                content_row["_content_type"] = col
                if struct_mod in IMAGE_MODALITIES and page_image_b64:
                    bbox = item.get("bbox_xyxy_norm")
                    if bbox and len(bbox) == 4:
                        cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_image_b64, bbox_xyxy_norm=bbox)
                        content_row["_image_b64"] = cropped_b64
                    else:
                        content_row["_image_b64"] = page_image_b64
                elif struct_mod in IMAGE_MODALITIES:
                    content_row["_image_b64"] = None
                new_rows.append(content_row)
                exploded_any = True

        # Preserve row if nothing was exploded (no text, no content).
        if not exploded_any:
            preserved = _deep_copy_row(row_dict)
            preserved["_embed_modality"] = text_mod
            preserved["_content_type"] = "text"
            if text_mod in IMAGE_MODALITIES:
                preserved["_image_b64"] = page_image_b64
            new_rows.append(preserved)

    return pd.DataFrame(new_rows).reset_index(drop=True)


def collapse_content_to_page_rows(
    batch_df: Any,
    *,
    text_column: str = "text",
    content_columns: Sequence[str] = _CONTENT_COLUMNS,
    modality: str = "text",
) -> Any:
    """Collapse each page into a single row for page-level embedding.

    For each row in *batch_df*:
    - Concatenates the page text with all structured content text (tables,
      charts, infographics) into one combined string in ``text_column``.
    - Uses the **full page image** (no cropping) when *modality* requires
      images.
    - Tags the row with ``_embed_modality = modality``.

    This produces **one embedding per page**.

    Parameters
    ----------
    modality : str
        Embedding modality for every row (default ``"text"``).
    """
    if not isinstance(batch_df, pd.DataFrame) or batch_df.empty:
        return batch_df

    batch_df = batch_df.copy()

    # Concatenate all text content per page using the existing helper.
    batch_df[text_column] = batch_df.apply(
        lambda row: _combine_text_with_content(row, text_column, content_columns),
        axis=1,
    )

    # Full page image (no cropping) for image modalities.
    if modality in IMAGE_MODALITIES:
        if "page_image" in batch_df.columns:
            batch_df["_image_b64"] = batch_df["page_image"].apply(
                lambda pi: pi.get("image_b64") if isinstance(pi, dict) else None
            )
        else:
            batch_df["_image_b64"] = None

    batch_df["_embed_modality"] = modality
    return batch_df


def _embed_group(
    group_df: pd.DataFrame,
    *,
    group_modality: str,
    model: Any,
    endpoint: Optional[str],
    api_key: Optional[str],
    text_column: str,
    inference_batch_size: int,
    output_column: str,
    resolved_model_name: str,
) -> pd.DataFrame:
    """Embed a single modality group via ``create_text_embeddings_for_df``.

    This is an internal helper called by ``embed_text_main_text_embed``
    once per distinct ``_embed_modality`` value found in the batch.
    """
    _embed = None
    _multimodal_embedder = None

    if endpoint is None and model is not None:
        if group_modality in IMAGE_MODALITIES:
            _multimodal_embedder = model
        else:
            _skip_prefix = hasattr(model, "embed_queries")

            def _embed(texts: Sequence[str]) -> Sequence[Sequence[float]]:  # noqa: F811
                batch = texts if _skip_prefix else [f"passage: {t}" for t in texts]
                vecs = model.embed(batch, batch_size=int(inference_batch_size))
                tolist = getattr(vecs, "tolist", None)
                if callable(tolist):
                    return tolist()
                return vecs  # type: ignore[return-value]

    # For remote image/text_image embedding, cap the HTTP batch size to avoid
    # sending huge JSON payloads (each item includes a base64 page image).
    _DEFAULT_REMOTE_IMAGE_BATCH_SIZE = 4
    effective_batch_size = inference_batch_size
    if endpoint is not None and group_modality in IMAGE_MODALITIES:
        effective_batch_size = min(inference_batch_size, _DEFAULT_REMOTE_IMAGE_BATCH_SIZE)

    cfg = TextEmbeddingConfig(
        text_column=str(text_column),
        output_payload_column=str(output_column) if output_column else None,
        write_embedding_to_metadata=True,
        metadata_column="metadata",
        batch_size=int(effective_batch_size),
        encoding_format="float",
        input_type="passage",
        truncate="END",
        dimensions=None,
        embedding_nim_endpoint=endpoint or "http://localhost:8012/v1",
        embedding_model=resolved_model_name or "nvidia/llama-nemotron-embed-1b-v2",
        embed_modality=group_modality,
    )

    out_df, _ = create_text_embeddings_for_df(
        group_df,
        task_config={
            "api_key": api_key,
            "embedder": _embed,
            "multimodal_embedder": _multimodal_embedder,
            "endpoint_url": endpoint,
            "local_batch_size": int(inference_batch_size),
        },
        transform_config=cfg,
    )
    return out_df


def embed_text_main_text_embed(
    batch_df: Any,
    *,
    model: Any = None,
    # Keep compatibility with previous `embed_text_1b_v2` signature so `.embed(...)` kwargs
    # don't need to change for inprocess mode.
    model_name: Optional[str] = None,
    embedding_endpoint: Optional[str] = None,
    embed_invoke_url: Optional[str] = None,
    api_key: Optional[str] = None,
    text_column: str = "text",
    inference_batch_size: int = 16,
    output_column: str = "text_embeddings_1b_v2",
    embedding_dim_column: str = "text_embeddings_1b_v2_dim",
    has_embedding_column: str = "text_embeddings_1b_v2_has_embedding",
    embed_modality: str = "text",
    **_: Any,
) -> Any:
    """
    Inprocess embedding task implemented via `nemo_retriever.text_embed.main_text_embed`.

    This is a thin adapter that preserves the old output columns:
    - `output_column` (payload dict with embedding)
    - `embedding_dim_column` (int)
    - `has_embedding_column` (bool)

    It also writes `metadata.embedding` for compatibility with downstream uploaders.

    When *embedding_endpoint* is set (e.g. ``"http://embedding:8000/v1"``), a remote
    NIM endpoint is used instead of the local HF model.  The NIM handles prefixing
    (``input_type="passage"``) server-side.

    Rows are grouped by their ``_embed_modality`` column (populated by
    ``explode_content_to_rows``) so that different element types can be
    embedded with different methods (e.g. text-only for page text,
    text+image for tables).  If no ``_embed_modality`` column exists,
    *embed_modality* is used uniformly.
    """
    if not isinstance(batch_df, pd.DataFrame):
        raise NotImplementedError("embed_text_main_text_embed currently only supports pandas.DataFrame input.")
    if inference_batch_size <= 0:
        raise ValueError("inference_batch_size must be > 0")

    # Resolve endpoint: strip whitespace, treat empty string as None.
    _endpoint = (embedding_endpoint or embed_invoke_url or "").strip() or None

    if _endpoint is None and model is None:
        raise ValueError("Either a local model or an embedding_endpoint must be provided.")

    # Resolve NIM aliases to the actual HF model ID.
    from nemo_retriever.model import resolve_embed_model

    _resolved_model_name = resolve_embed_model(model_name)

    # Determine per-row modalities.  If _embed_modality column exists (set by
    # explode_content_to_rows), group by it; otherwise fall back to the single
    # embed_modality parameter.
    has_per_row_modality = "_embed_modality" in batch_df.columns
    if has_per_row_modality:
        modalities = batch_df["_embed_modality"].fillna(embed_modality).unique().tolist()
    else:
        modalities = [embed_modality]

    try:
        if len(modalities) == 1:
            # Fast path: all rows share the same modality — process in one shot.
            out_df = _embed_group(
                batch_df,
                group_modality=modalities[0],
                model=model,
                endpoint=_endpoint,
                api_key=api_key,
                text_column=text_column,
                inference_batch_size=inference_batch_size,
                output_column=output_column,
                resolved_model_name=_resolved_model_name,
            )
        else:
            # Multiple modalities: group, embed each, reassemble in original order.
            parts: List[pd.DataFrame] = []
            for mod in modalities:
                mask = batch_df["_embed_modality"] == mod
                group_df = batch_df.loc[mask]
                if group_df.empty:
                    continue
                part = _embed_group(
                    group_df,
                    group_modality=mod,
                    model=model,
                    endpoint=_endpoint,
                    api_key=api_key,
                    text_column=text_column,
                    inference_batch_size=inference_batch_size,
                    output_column=output_column,
                    resolved_model_name=_resolved_model_name,
                )
                parts.append(part)
            out_df = pd.concat(parts).sort_index()

    except BaseException as e:
        # Fail-soft: preserve batch shape and set an error payload per-row.
        import traceback as _tb

        print(f"Warning: embedding failed: {type(e).__name__}: {e}")
        _tb.print_exc()
        err_payload = {"embedding": None, "error": {"stage": "embed", "type": e.__class__.__name__, "message": str(e)}}
        out_df = batch_df.copy()
        if output_column:
            out_df[output_column] = [err_payload for _ in range(len(out_df.index))]
        out_df[embedding_dim_column] = [0 for _ in range(len(out_df.index))]
        out_df[has_embedding_column] = [False for _ in range(len(out_df.index))]
        out_df["_contains_embeddings"] = [False for _ in range(len(out_df.index))]
        return out_df

    # Add backwards-compatible dim/has columns (these were produced by `embed_text_1b_v2`).
    if embedding_dim_column:

        def _dim(row: pd.Series) -> int:
            md = row.get("metadata")
            if isinstance(md, dict):
                emb = md.get("embedding")
                if isinstance(emb, list):
                    return int(len(emb))
            payload = row.get(output_column) if output_column else None
            if isinstance(payload, dict) and isinstance(payload.get("embedding"), list):
                return int(len(payload.get("embedding") or []))
            return 0

        out_df[embedding_dim_column] = out_df.apply(_dim, axis=1)
    else:
        out_df[embedding_dim_column] = [0 for _ in range(len(out_df.index))]

    out_df[has_embedding_column] = [bool(int(d) > 0) for d in out_df[embedding_dim_column].tolist()]

    # Drop helper columns after embedding to free memory.
    for col in ("_image_b64", "_embed_modality"):
        if col in out_df.columns:
            out_df = out_df.drop(columns=[col])

    return out_df


def _to_jsonable(obj: Any) -> Any:
    """
    Best-effort conversion of arbitrary objects into JSON-serializable types.
    """
    if obj is None or isinstance(obj, (str, int, float, bool)):
        return obj

    if isinstance(obj, dict):
        return {str(k): _to_jsonable(v) for k, v in obj.items()}

    if isinstance(obj, (list, tuple)):
        return [_to_jsonable(v) for v in obj]

    # Common numeric scalar types (numpy, pandas).
    item = getattr(obj, "item", None)
    if callable(item):
        try:
            return _to_jsonable(item())
        except Exception:
            pass

    # Common array/tensor types (numpy, torch).
    tolist = getattr(obj, "tolist", None)
    if callable(tolist):
        try:
            return _to_jsonable(tolist())
        except Exception:
            pass

    # Pandas timestamps.
    isoformat = getattr(obj, "isoformat", None)
    if callable(isoformat):
        try:
            return str(isoformat())
        except Exception:
            pass

    return str(obj)


def _atomic_write_json(path: str, payload: Any) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
        f.flush()
        try:
            os.fsync(f.fileno())
        except Exception:
            pass
    os.replace(tmp, path)


def _safe_stem(name: str) -> str:
    s = str(name or "").strip() or "document"
    s = os.path.splitext(os.path.basename(s))[0] or "document"
    # Keep it filesystem-friendly.
    s = re.sub(r"[^A-Za-z0-9._-]+", "_", s)
    return s[:160] if len(s) > 160 else s


def pages_df_from_pdf_bytes(pdf_bytes: Union[bytes, bytearray], source_path: str) -> pd.DataFrame:
    """
    Build a per-page DataFrame from raw PDF bytes (same schema as pdf_path_to_pages_df).

    Used by the online ingest mode to run the same pipeline on document bytes
    received via REST. Columns: bytes, path, page_number.
    """
    pages = _split_pdf_to_single_page_bytes(pdf_bytes)
    out_rows = [{"bytes": b, "path": source_path, "page_number": i + 1} for i, b in enumerate(pages)]
    return pd.DataFrame(out_rows)


def run_pipeline_tasks_on_df(
    initial_df: pd.DataFrame,
    per_doc_tasks: Sequence[Tuple[Callable[..., Any], Dict[str, Any]]],
    post_tasks: Optional[Sequence[Tuple[Callable[..., Any], Dict[str, Any]]]] = None,
) -> Tuple[Any, List[Dict[str, Any]]]:
    """
    Run the inprocess pipeline task chain on a single document DataFrame.

    Returns (final_result, metrics) where metrics is a list of
    {"stage": str, "duration_sec": float} for each stage. Used by both
    InProcessIngestor.ingest() and the online Ray Serve deployment.
    """
    import time

    metrics: List[Dict[str, Any]] = []
    current: Any = initial_df
    for func, kwargs in per_doc_tasks:
        t0 = time.perf_counter()
        if func is pdf_extraction:
            current = func(pdf_binary=current, **kwargs)
        else:
            current = func(current, **kwargs)
        metrics.append({"stage": getattr(func, "__name__", "unknown"), "duration_sec": time.perf_counter() - t0})

    if post_tasks and current is not None:
        combined = current if isinstance(current, pd.DataFrame) else pd.concat(current, ignore_index=True)
        for func, kwargs in post_tasks:
            t0 = time.perf_counter()
            combined = func(combined, **kwargs)
            metrics.append({"stage": getattr(func, "__name__", "unknown"), "duration_sec": time.perf_counter() - t0})
        return combined, metrics
    return current, metrics


def save_dataframe_to_disk_json(df: Any, *, output_directory: str) -> Any:
    """
    Pipeline task: persist a pandas DataFrame as a timestamped JSON file.

    Returns the input unchanged so it can stay in the pipeline.
    """
    if not isinstance(df, pd.DataFrame):
        raise TypeError(f"save_dataframe_to_disk_json expects pandas.DataFrame, got {type(df)!r}")
    if not isinstance(output_directory, str) or not output_directory.strip():
        raise ValueError("output_directory must be a non-empty string")

    out_dir = os.path.abspath(output_directory)
    os.makedirs(out_dir, exist_ok=True)

    # Try to derive a human-friendly filename prefix from the source PDF path.
    source_path = None
    try:
        if "metadata" in df.columns and len(df.index) > 0:
            md = df.iloc[0]["metadata"]
            if isinstance(md, dict):
                source_path = md.get("source_path")
    except Exception:
        source_path = None

    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    stem = _safe_stem(str(source_path) if source_path else "results")
    uniq = uuid.uuid4().hex[:8]
    out_path = os.path.join(out_dir, f"{stem}.{ts}.{uniq}.json")

    payload = {
        "schema_version": 1,
        "run_mode": "inprocess",
        "created_at_utc": ts,
        "source_path": str(source_path) if source_path else None,
        "rows": int(len(df.index)),
        "columns": [str(c) for c in df.columns.tolist()],
        "records": _to_jsonable(df.to_dict(orient="records")),
    }

    _atomic_write_json(out_path, payload)
    return df


from nemo_retriever.ingest_modes.lancedb_utils import (
    build_lancedb_rows,
    create_or_append_lancedb_table,
    infer_vector_dim,
    lancedb_schema,
)


def upload_embeddings_to_lancedb_inprocess(
    df: Any,
    *,
    lancedb_uri: str = "lancedb",
    table_name: str = "nv-ingest",
    overwrite: bool = True,
    create_index: bool = True,
    index_type: str = "IVF_HNSW_SQ",
    metric: str = "l2",
    num_partitions: int = 16,
    num_sub_vectors: int = 256,
    embedding_column: str = "text_embeddings_1b_v2",
    embedding_key: str = "embedding",
    include_text: bool = True,
    text_column: str = "text",
    hybrid: bool = False,
    fts_language: str = "English",
    **_ignored: Any,
) -> Any:
    """
    Pipeline task: upload embeddings from a pandas DataFrame into LanceDB.

    The embedding is sourced from:
    - `metadata.embedding` if present, else
    - `embedding_column` payloads like those produced by `embed_text_1b_v2`

    Parameters (all can be passed via `.vdb_upload(...)` kwargs):
    - **lancedb_uri**: LanceDB URI (directory path), default `"lancedb"`
    - **table_name**: table name, default `"nv-ingest"`
    - **overwrite**: overwrite table (True) or append (False)
    - **create_index**: create vector index after upload
    - **index_type**: LanceDB index type (e.g. `"IVF_HNSW_SQ"`)
    - **metric**: distance metric (e.g. `"l2"`, `"cosine"`)
    - **num_partitions**: index partitions
    - **num_sub_vectors**: index sub-vectors
    - **embedding_column**: column name containing embedding payload dicts
    - **embedding_key**: key inside payload dict to read embedding list from
    - **include_text**: if True, store `text_column` content alongside the vector
    - **text_column**: column to read text from when `include_text=True`
    - **hybrid**: if True, additionally build an FTS index on `text`
    - **fts_language**: language passed to LanceDB FTS index creation

    Returns the input DataFrame unchanged (so it can remain in the pipeline).
    """
    if not isinstance(df, pd.DataFrame):
        raise TypeError(f"upload_embeddings_to_lancedb_inprocess expects pandas.DataFrame, got {type(df)!r}")

    rows = build_lancedb_rows(
        df,
        embedding_column=str(embedding_column),
        embedding_key=str(embedding_key),
        text_column=str(text_column),
        include_text=bool(include_text),
    )

    if not rows:
        print("No embeddings found to upload to LanceDB (no rows had embeddings).")
        return df

    dim = infer_vector_dim(rows)
    if dim <= 0:
        raise ValueError("Failed to infer embedding dimension from DataFrame rows.")

    try:
        import lancedb  # type: ignore
    except Exception as e:
        raise RuntimeError(
            "LanceDB upload requested but dependencies are missing. Install `lancedb` and `pyarrow`."
        ) from e

    db = lancedb.connect(uri=str(lancedb_uri))
    schema = lancedb_schema(dim)
    table = create_or_append_lancedb_table(db, str(table_name), rows, schema, overwrite=overwrite)

    if create_index:
        # LanceDB IVF-based indexes train k-means with K=num_partitions. K must be < N vectors.
        n_vecs = int(len(rows))
        if n_vecs < 2:
            print("Skipping LanceDB index creation (not enough vectors).")
        else:
            k = int(num_partitions)
            if k >= n_vecs:
                k = max(1, n_vecs - 1)
            try:
                table.create_index(
                    index_type=str(index_type),
                    metric=str(metric),
                    num_partitions=int(k),
                    num_sub_vectors=int(num_sub_vectors),
                    vector_column_name="vector",
                )
            except TypeError:
                # Older/newer LanceDB versions may have different signatures; fall back to minimal call.
                table.create_index(vector_column_name="vector")
            except Exception as e:
                # Don't fail ingestion due to index training; users can rebuild later with more data.
                print(f"Warning: failed to create LanceDB index (continuing without index): {e}")

    if hybrid:
        # Build text index for hybrid retrieval when text is present.
        try:
            table.create_fts_index("text", language=str(fts_language))
        except Exception as e:
            # Keep ingestion resilient: vector data was already written.
            print(f"Warning: failed to create LanceDB FTS index (continuing without FTS): {e}")

    print(f"Wrote {len(rows)} rows to LanceDB uri={lancedb_uri!r} table={table_name!r}")
    return df


def pdf_to_pages_df(path: str) -> pd.DataFrame:
    """
    Convert a document at *path* into a DataFrame where each row
    contains a *single-page* PDF's raw bytes.

    For ``.docx`` / ``.pptx`` files the document is first converted
    to PDF via LibreOffice, then split into pages.  The original
    *path* is preserved so downstream metadata tracks the source file.

    Columns:
    - bytes: single-page PDF bytes
    - path: original input path
    - page_number: 1-indexed page number
    """
    if pdfium is None:  # pragma: no cover
        raise ImportError("pypdfium2 is required for inprocess ingestion.") from _PDFIUM_IMPORT_ERROR

    abs_path = os.path.abspath(path)
    ext = os.path.splitext(abs_path)[1].lower()
    out_rows: list[dict[str, Any]] = []
    doc = None
    try:
        if ext in SUPPORTED_EXTENSIONS and ext != ".pdf":
            # Convert DOCX/PPTX to PDF bytes first.
            with open(abs_path, "rb") as f:
                file_bytes = f.read()
            pdf_bytes = convert_to_pdf_bytes(file_bytes, ext)
            doc = pdfium.PdfDocument(BytesIO(pdf_bytes))
        else:
            doc = pdfium.PdfDocument(abs_path)

        for page_idx in range(len(doc)):
            single = pdfium.PdfDocument.new()
            try:
                single.import_pages(doc, pages=[page_idx])
                buf = BytesIO()
                single.save(buf)
                out_rows.append(
                    {
                        "bytes": buf.getvalue(),
                        "path": abs_path,
                        "page_number": page_idx + 1,
                    }
                )
            finally:
                try:
                    single.close()
                except Exception:
                    pass
    except BaseException as e:
        # Preserve shape expected downstream (pdf_extraction emits error
        # records per-row, so we return a single row to trigger that).
        out_rows.append({"bytes": b"", "path": abs_path, "page_number": 0, "error": str(e)})
    finally:
        try:
            if doc is not None:
                doc.close()
        except Exception:
            pass

    return pd.DataFrame(out_rows)


def _iter_page_chunks(path: str, chunk_size: int = 32) -> Iterator[pd.DataFrame]:
    """Yield DataFrames of *chunk_size* pages from a document.

    Reuses the same pdfium page-splitting logic as :func:`pdf_to_pages_df`
    but yields incrementally so that downstream work can begin before the
    entire document is split.  At most *chunk_size* single-page PDF blobs
    are in memory per yield.

    Handles DOCX/PPTX conversion (same as ``pdf_to_pages_df``).
    """
    if pdfium is None:  # pragma: no cover
        raise ImportError("pypdfium2 is required for inprocess ingestion.") from _PDFIUM_IMPORT_ERROR

    abs_path = os.path.abspath(path)
    ext = os.path.splitext(abs_path)[1].lower()
    doc = None
    try:
        if ext in SUPPORTED_EXTENSIONS and ext != ".pdf":
            with open(abs_path, "rb") as f:
                pdf_bytes = convert_to_pdf_bytes(f.read(), ext)
            doc = pdfium.PdfDocument(BytesIO(pdf_bytes))
        else:
            doc = pdfium.PdfDocument(abs_path)

        chunk: list[dict] = []
        for page_idx in range(len(doc)):
            single = pdfium.PdfDocument.new()
            try:
                single.import_pages(doc, pages=[page_idx])
                buf = BytesIO()
                single.save(buf)
                chunk.append({"bytes": buf.getvalue(), "path": abs_path, "page_number": page_idx + 1})
            finally:
                try:
                    single.close()
                except Exception:
                    pass
            if len(chunk) >= chunk_size:
                yield pd.DataFrame(chunk)
                chunk = []
        if chunk:
            yield pd.DataFrame(chunk)
    except BaseException as e:
        yield pd.DataFrame([{"bytes": b"", "path": abs_path, "page_number": 0, "error": str(e)}])
    finally:
        if doc is not None:
            try:
                doc.close()
            except Exception:
                pass


def _process_doc_cpu(doc_path: str, cpu_tasks: list) -> pd.DataFrame:
    """Worker function for ProcessPoolExecutor.

    Runs ``pdf_to_pages_df`` followed by all CPU-bound pipeline tasks on a
    single document.  All arguments are picklable (strings, module-level
    functions, and dicts of simple types).
    """
    try:
        current: Any = pdf_to_pages_df(doc_path)
        for func, kwargs in cpu_tasks:
            if func is pdf_extraction:
                current = func(pdf_binary=current, **kwargs)
            else:
                current = func(current, **kwargs)
        return current
    except Exception as e:
        # Return a minimal error DataFrame so one failed document does not
        # halt the pipeline.
        return pd.DataFrame([{"bytes": b"", "path": doc_path, "page_number": 0, "error": f"{type(e).__name__}: {e}"}])


def _process_chunk_cpu(chunk_df: pd.DataFrame, cpu_tasks: list) -> pd.DataFrame:
    """Worker function for ProcessPoolExecutor — page-chunk variant.

    Runs CPU-bound pipeline tasks on a pre-split page-chunk DataFrame.
    Similar to :func:`_process_doc_cpu` but skips ``pdf_to_pages_df``
    since the page splitting has already been done by the caller.
    """
    try:
        current: Any = chunk_df
        for func, kwargs in cpu_tasks:
            if func is pdf_extraction:
                current = func(pdf_binary=current, **kwargs)
            else:
                current = func(current, **kwargs)
        return current
    except Exception as e:
        return pd.DataFrame([{"bytes": b"", "path": "", "page_number": 0, "error": f"{type(e).__name__}: {e}"}])


def _collect_summary_from_df(df: pd.DataFrame) -> dict:
    """Compute detection summary from a result DataFrame."""
    from nemo_retriever.utils.detection_summary import collect_detection_summary_from_df

    return collect_detection_summary_from_df(df)


def _print_ingest_summary(results: list, elapsed_s: float) -> None:
    """Print end-of-ingest summary matching batch pipeline output format."""
    from nemo_retriever.utils.detection_summary import print_detection_summary

    dfs = [r for r in results if isinstance(r, pd.DataFrame) and not r.empty]
    if not dfs:
        print(f"\nIngest time: {elapsed_s:.2f}s (no documents processed)")
        return

    combined = pd.concat(dfs, ignore_index=True) if len(dfs) > 1 else dfs[0]
    summary = _collect_summary_from_df(combined)
    print_detection_summary(summary)

    pages = summary["pages_seen"]
    if elapsed_s > 0 and pages > 0:
        pps = pages / elapsed_s
        print(f"Pages processed: {pages}")
        print(f"Pages/sec: {pps:.2f}")
    else:
        print(f"\nIngest time: {elapsed_s:.2f}s")


class InProcessIngestor(Ingestor):
    RUN_MODE = "inprocess"

    def __init__(self, documents: Optional[List[str]] = None) -> None:
        super().__init__(documents=documents)

        # Keep backwards-compatibility with code that inspects `Ingestor._documents`
        # by ensuring both names refer to the same list.
        self._input_documents: List[str] = self._documents

        # Builder-style configuration recorded for later execution (TBD).
        self._tasks: List[tuple[Callable[..., Any], dict[str, Any]]] = []

        # Pipeline type: "pdf" (extract), "txt" (extract_txt), "html" (extract_html), or "image" (extract_image_files).
        self._pipeline_type: Literal["pdf", "txt", "html", "image"] = "pdf"
        self._extract_txt_kwargs: Dict[str, Any] = {}
        self._extract_html_kwargs: Dict[str, Any] = {}
        self._caption_enabled: bool = False

    def files(self, documents: Union[str, List[str]]) -> "InProcessIngestor":
        """
        Add local files for in-process execution.

        Mirrors `BatchIngestor.files()` path/glob resolution, but does not create
        a Ray Data dataset.
        """
        if isinstance(documents, str):
            documents = [documents]

        for pattern in documents:
            if not isinstance(pattern, str) or not pattern:
                raise ValueError(f"Invalid document pattern: {pattern!r}")

            # Expand globs (supports ** when recursive=True).
            matches = glob.glob(pattern, recursive=True)
            if matches:
                files = [os.path.abspath(p) for p in matches if os.path.isfile(p)]
                if not files:
                    raise FileNotFoundError(f"Pattern resolved, but no files found: {pattern!r}")
                self._input_documents.extend(files)
                continue

            # No glob matches: treat as explicit path.
            if os.path.isfile(pattern):
                self._input_documents.append(os.path.abspath(pattern))
                continue

            raise FileNotFoundError(f"No local files found for: {pattern!r}")

        return self

    def extract(self, params: ExtractParams | None = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Configure extraction for in-process execution (skeleton).

        TODO: implement actual local extraction (PDF split/extract, page elements,
        table structure, chart detection, etc.). For now this records the
        configuration so call sites can be wired up and chained.
        """
        # NOTE: `kwargs` passed to `.extract()` are intended primarily for PDF extraction
        # (e.g. `extract_text`, `dpi`, etc). Downstream model stages do NOT necessarily
        # accept the same keyword arguments. Keep per-stage kwargs isolated.
        if self._input_documents and all(f.lower().endswith(".txt") for f in self._input_documents):
            txt_params = TextChunkParams()
            return self.extract_txt(params=txt_params)
        if self._input_documents and all(f.lower().endswith(".html") for f in self._input_documents):
            html_params = HtmlChunkParams()
            return self.extract_html(params=html_params)
        if self._input_documents and all(
            os.path.splitext(f)[1].lower() in SUPPORTED_IMAGE_EXTENSIONS for f in self._input_documents
        ):
            return self.extract_image_files(params=params, **kwargs)
        resolved = _coerce_params(params, ExtractParams, kwargs)
        if (
            any(
                (
                    resolved.invoke_url,
                    resolved.page_elements_invoke_url,
                    resolved.ocr_invoke_url,
                    resolved.graphic_elements_invoke_url,
                    resolved.table_structure_invoke_url,
                )
            )
            and not resolved.api_key
        ):
            resolved = resolved.model_copy(update={"api_key": resolve_remote_api_key()})
        kwargs = resolved.model_dump(mode="python")
        use_nemotron_parse_only = kwargs.get("method") == "nemotron_parse"
        extract_kwargs = dict(kwargs)
        # Downstream in-process stages (page elements / table / chart / infographic) assume
        # `page_image.image_b64` exists. Ensure PDF extraction emits a page image unless
        # the caller explicitly set it to False.
        if extract_kwargs.get("extract_page_as_image") is None:
            if any(
                extract_kwargs.get(k) is True
                for k in ("extract_text", "extract_images", "extract_tables", "extract_charts", "extract_infographics")
            ):
                extract_kwargs["extract_page_as_image"] = True
        self._pipeline_type = "pdf"
        self._tasks.append((pdf_extraction, extract_kwargs))

        self._append_detection_tasks(kwargs, use_nemotron_parse_only=use_nemotron_parse_only)

        return self

    def _append_detection_tasks(
        self,
        kwargs: dict[str, Any],
        *,
        use_nemotron_parse_only: bool = False,
    ) -> None:
        """Append page-element detection, OCR, table/chart/infographic tasks.

        Shared by ``extract()`` (PDF) and ``extract_image_files()`` (standalone images).
        """
        # Common, optional knobs shared by our detect_* helpers.
        detect_passthrough_keys = {
            "inference_batch_size",
            "output_column",
            "num_detections_column",
            "counts_by_label_column",
        }

        def _stage_remote_kwargs(stage_name: str) -> dict[str, Any]:
            stage_prefix = f"{stage_name}_"
            out: dict[str, Any] = {}
            invoke_url = kwargs.get(f"{stage_prefix}invoke_url", kwargs.get("invoke_url"))
            if invoke_url:
                out["invoke_url"] = invoke_url
            api_key = kwargs.get(f"{stage_prefix}api_key", kwargs.get("api_key"))
            if api_key:
                out["api_key"] = api_key
            timeout = kwargs.get(f"{stage_prefix}request_timeout_s", kwargs.get("request_timeout_s"))
            if timeout is not None:
                out["request_timeout_s"] = timeout
            for k in ("remote_max_pool_workers", "remote_max_retries", "remote_max_429_retries"):
                stage_key = f"{stage_prefix}{k}"
                if stage_key in kwargs:
                    out[k] = kwargs[stage_key]
                elif k in kwargs:
                    out[k] = kwargs[k]
            return out

        def _detect_kwargs_with_model(model_obj: Any, *, stage_name: str, allow_remote: bool) -> dict[str, Any]:
            d: dict[str, Any] = {"model": model_obj}
            for k in detect_passthrough_keys:
                if k in kwargs:
                    d[k] = kwargs[k]
            if allow_remote:
                d.update(_stage_remote_kwargs(stage_name))
            return d

        if use_nemotron_parse_only:
            parse_flags: dict[str, Any] = {}
            if kwargs.get("extract_tables") is True:
                parse_flags["extract_tables"] = True
            if kwargs.get("extract_charts") is True:
                parse_flags["extract_charts"] = True
            if kwargs.get("extract_infographics") is True:
                parse_flags["extract_infographics"] = True
            parse_flags.update(_stage_remote_kwargs("nemotron_parse"))
            parse_invoke_url = kwargs.get(
                "nemotron_parse_invoke_url", kwargs.get("ocr_invoke_url", kwargs.get("invoke_url", ""))
            )
            if parse_invoke_url:
                parse_flags["invoke_url"] = parse_invoke_url
                self._tasks.append((nemotron_parse_page_elements, {"model": None, **parse_flags}))
            else:
                self._tasks.append((nemotron_parse_page_elements, {"model": NemotronParseV12(), **parse_flags}))
        else:
            # Page-elements detection is needed for tables/charts/infographics,
            # and for text only when the method requires OCR (not plain pdfium).
            _method = kwargs.get("method", "pdfium")
            _need_pe_for_text = kwargs.get("extract_text") is True and _method in ("pdfium_hybrid", "ocr")
            if _need_pe_for_text or any(
                kwargs.get(k) is True for k in ("extract_tables", "extract_charts", "extract_infographics")
            ):
                print("Adding page elements task")
                pe_invoke_url = kwargs.get("page_elements_invoke_url", kwargs.get("invoke_url", ""))
                pe_model = None if pe_invoke_url else NemotronPageElementsV3()
                self._tasks.append(
                    (
                        detect_page_elements_v3,
                        _detect_kwargs_with_model(
                            pe_model,
                            stage_name="page_elements",
                            allow_remote=True,
                        ),
                    )
                )

            use_graphic_elements = bool(kwargs.get("use_graphic_elements", False))

            # When use_graphic_elements is True, charts go through the
            # combined graphic-elements + OCR stage instead of OCR-only.
            if use_graphic_elements and kwargs.get("extract_charts") is True:
                print("Adding graphic-elements+OCR extraction task")
                ge_invoke_url = kwargs.get("graphic_elements_invoke_url", "")
                ocr_invoke_url = kwargs.get("ocr_invoke_url", kwargs.get("invoke_url", ""))
                ocr_model_dir = (
                    kwargs.get("ocr_model_dir")
                    or os.environ.get("RETRIEVER_NEMOTRON_OCR_MODEL_DIR", "").strip()
                    or os.environ.get("NEMOTRON_OCR_MODEL_DIR", "").strip()
                    or os.environ.get("NEMOTRON_OCR_V1_MODEL_DIR", "").strip()
                )

                ge_ocr_kwargs: dict[str, Any] = {}
                if ge_invoke_url:
                    ge_ocr_kwargs["graphic_elements_invoke_url"] = ge_invoke_url
                    ge_ocr_kwargs["graphic_elements_model"] = None
                else:
                    from nemo_retriever.model.local import NemotronGraphicElementsV1

                    ge_ocr_kwargs["graphic_elements_model"] = NemotronGraphicElementsV1()

                if ocr_invoke_url:
                    ge_ocr_kwargs["ocr_invoke_url"] = ocr_invoke_url
                    ge_ocr_kwargs["ocr_model"] = None
                else:
                    ocr_model = NemotronOCRV1(model_dir=str(ocr_model_dir)) if ocr_model_dir else NemotronOCRV1()
                    ge_ocr_kwargs["ocr_model"] = ocr_model

                ge_ocr_kwargs.update(_stage_remote_kwargs("ocr"))
                self._tasks.append((graphic_elements_ocr_page_elements, ge_ocr_kwargs))

            use_table_structure = bool(kwargs.get("use_table_structure", False))
            from nemo_retriever.application.pipeline.build_plan import validate_table_structure_flags

            validate_table_structure_flags(
                use_table_structure, str(kwargs.get("table_output_format", "pseudo_markdown"))
            )

            # When use_table_structure is True, tables go through
            # the combined table-structure + OCR stage instead of OCR-only.
            if use_table_structure and kwargs.get("extract_tables") is True:
                print("Adding table-structure+OCR extraction task")
                ts_invoke_url = kwargs.get("table_structure_invoke_url", "")
                ocr_invoke_url = kwargs.get("ocr_invoke_url", kwargs.get("invoke_url", ""))
                ocr_model_dir = (
                    kwargs.get("ocr_model_dir")
                    or os.environ.get("RETRIEVER_NEMOTRON_OCR_MODEL_DIR", "").strip()
                    or os.environ.get("NEMOTRON_OCR_MODEL_DIR", "").strip()
                    or os.environ.get("NEMOTRON_OCR_V1_MODEL_DIR", "").strip()
                )

                ts_ocr_kwargs: dict[str, Any] = {}
                if ts_invoke_url:
                    ts_ocr_kwargs["table_structure_invoke_url"] = ts_invoke_url
                    ts_ocr_kwargs["table_structure_model"] = None
                else:
                    from nemo_retriever.model.local import NemotronTableStructureV1

                    ts_ocr_kwargs["table_structure_model"] = NemotronTableStructureV1()

                if ocr_invoke_url:
                    ts_ocr_kwargs["ocr_invoke_url"] = ocr_invoke_url
                    ts_ocr_kwargs["ocr_model"] = None
                else:
                    ocr_model = NemotronOCRV1(model_dir=str(ocr_model_dir)) if ocr_model_dir else NemotronOCRV1()
                    ts_ocr_kwargs["ocr_model"] = ocr_model

                ts_ocr_kwargs.update(_stage_remote_kwargs("ocr"))
                self._tasks.append((table_structure_ocr_page_elements, ts_ocr_kwargs))

            # OCR-based extraction for tables/charts/infographics.
            # When use_graphic_elements is True, charts are handled above;
            # when use_table_structure is True, tables are handled above.
            ocr_flags = {}
            method = kwargs.get("method", "pdfium")
            if method in ("pdfium_hybrid", "ocr") and kwargs.get("extract_text") is True:
                ocr_flags["extract_text"] = True
            if kwargs.get("extract_tables") is True and not use_table_structure:
                ocr_flags["extract_tables"] = True
            if kwargs.get("extract_charts") is True and not use_graphic_elements:
                ocr_flags["extract_charts"] = True
            if kwargs.get("extract_infographics") is True:
                ocr_flags["extract_infographics"] = True
            ocr_flags.update(_stage_remote_kwargs("ocr"))

            if ocr_flags:
                print("Adding OCR extraction task")
                ocr_invoke_url = kwargs.get("ocr_invoke_url", kwargs.get("invoke_url", ""))
                if ocr_invoke_url:
                    self._tasks.append((ocr_page_elements, {"model": None, **ocr_flags}))
                else:
                    ocr_model_dir = (
                        kwargs.get("ocr_model_dir")
                        or os.environ.get("RETRIEVER_NEMOTRON_OCR_MODEL_DIR", "").strip()
                        or os.environ.get("NEMOTRON_OCR_MODEL_DIR", "").strip()
                        or os.environ.get("NEMOTRON_OCR_V1_MODEL_DIR", "").strip()
                    )
                    model = NemotronOCRV1(model_dir=str(ocr_model_dir)) if ocr_model_dir else NemotronOCRV1()
                    self._tasks.append((ocr_page_elements, {"model": model, **ocr_flags}))

    def extract_image_files(self, params: ExtractParams | None = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Configure image ingestion: standalone image -> page DataFrame -> detection/OCR.

        Use with .files("*.png").extract_image_files(...).embed().vdb_upload().ingest().
        Do not call .extract() when using .extract_image_files().
        """
        resolved = _coerce_params(params, ExtractParams, kwargs)
        if (
            any(
                (
                    resolved.invoke_url,
                    resolved.page_elements_invoke_url,
                    resolved.ocr_invoke_url,
                    resolved.graphic_elements_invoke_url,
                    resolved.table_structure_invoke_url,
                )
            )
            and not resolved.api_key
        ):
            resolved = resolved.model_copy(update={"api_key": resolve_remote_api_key()})
        kwargs = resolved.model_dump(mode="python")
        use_nemotron_parse_only = kwargs.get("method") == "nemotron_parse"
        self._pipeline_type = "image"
        self._append_detection_tasks(kwargs, use_nemotron_parse_only=use_nemotron_parse_only)
        return self

    def split(self, params: TextChunkParams | None = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Re-chunk the ``text`` column by token count (post-extraction transform).

        Appends :func:`~nemo_retriever.txt.split.split_df` as a GPU-category
        task so it runs in sequence after extraction and before embedding.
        """
        from nemo_retriever.txt.split import split_df

        resolved = _coerce_params(params, TextChunkParams, kwargs)
        split_kwargs = resolved.model_dump(mode="python")
        split_kwargs.pop("encoding", None)
        self._tasks.append((split_df, split_kwargs))
        return self

    def extract_txt(self, params: TextChunkParams | None = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Configure txt ingestion: tokenizer-based chunking only (no PDF extraction).

        Use with .files("*.txt").extract_txt(...).embed().vdb_upload().ingest().
        Do not call .extract() when using .extract_txt().
        """
        from nemo_retriever.txt.ray_data import TxtSplitActor

        self._pipeline_type = "txt"
        resolved = _coerce_params(params, TextChunkParams, kwargs)
        self._extract_txt_kwargs = resolved.model_dump(mode="python")
        text_split = TxtSplitActor(params=TextChunkParams(**self._extract_txt_kwargs))
        self._tasks.append((text_split, {}))
        return self

    def extract_html(self, params: HtmlChunkParams | None = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Configure HTML ingestion: markitdown -> markdown -> tokenizer chunking (no PDF extraction).

        Use with .files("*.html").extract_html(...).embed().vdb_upload().ingest().
        Do not call .extract() when using .extract_html().
        """
        from nemo_retriever.html.ray_data import HtmlSplitActor

        self._pipeline_type = "html"
        resolved = _coerce_params(params, HtmlChunkParams, kwargs)
        self._extract_html_kwargs = resolved.model_dump(mode="python")
        html_split = HtmlSplitActor(
            params=HtmlChunkParams(**self._extract_html_kwargs),
        )
        self._tasks.append((html_split, {}))
        return self

    def extract_audio(
        self,
        params: AudioChunkParams | None = None,
        asr_params: ASRParams | None = None,
        **kwargs: Any,
    ) -> "InProcessIngestor":
        """
        Configure audio ingestion: media chunking (sync) -> ASR (Parakeet gRPC) -> same embed/vdb.

        Use with .files("mp3/*.mp3").extract_audio(...).embed().vdb_upload().ingest().
        Do not call .extract() when using .extract_audio(). ASR requires a remote Parakeet/Riva endpoint.
        """
        from nemo_retriever.audio.asr_actor import apply_asr_to_df

        self._pipeline_type = "audio"
        chunk_resolved = _coerce_params(params, AudioChunkParams, kwargs)
        asr_resolved = _coerce_params(asr_params, ASRParams, kwargs)
        self._extract_audio_chunk_kwargs = chunk_resolved.model_dump(mode="python")
        self._extract_audio_asr_kwargs = asr_resolved.model_dump(mode="python")
        self._tasks.append((apply_asr_to_df, {"asr_params": self._extract_audio_asr_kwargs}))
        return self

    def dedup(self, params: "DedupParams | None" = None, **kwargs: Any) -> "InProcessIngestor":
        """Remove duplicate and overlapping images before captioning."""

        resolved = _coerce_params(params, DedupParams, kwargs)
        self._tasks.append((dedup_images, resolved.model_dump(mode="python")))
        return self

    def caption(self, params: "CaptionParams | None" = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Configure image captioning via a local VLM model or remote endpoint.

        Sends cropped images (from the ``images`` column populated by
        ``extract(extract_images=True)``) to a VLM and writes the returned
        captions back as ``images[i]["text"]``.

        When ``endpoint_url`` is set, a remote NIM endpoint is used.
        Otherwise a local ``NemotronVLMCaptioner`` is loaded from HF.
        """
        from nemo_retriever.caption.caption import caption_images
        from nemo_retriever.params import CaptionParams

        resolved = _coerce_params(params, CaptionParams, kwargs)
        caption_kwargs = resolved.model_dump(mode="python")

        if resolved.endpoint_url:
            # Remote mode.
            if not resolved.api_key:
                caption_kwargs["api_key"] = resolve_remote_api_key()
        else:
            # Local mode: defer model creation so the VLM is loaded lazily
            # on the device specified by CaptionParams.device.
            if not resolved.device:
                import warnings

                warnings.warn(
                    "No caption device specified. The VLM will load on cuda:0, which "
                    "may conflict with other models. Use device='cuda:1' (or "
                    "--caption-device from the CLI) to place the captioner on a "
                    "separate GPU.",
                    stacklevel=2,
                )
            caption_kwargs["model"] = None

        self._caption_enabled = True
        self._tasks.append((caption_images, caption_kwargs))
        return self

    def embed(self, params: EmbedParams | None = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Configure embedding for in-process execution.

        This records an embedding task so call sites can chain `.embed(...)`
        after `.extract(...)`, `.extract_txt()`, or `.extract_html()`.

        When ``embedding_endpoint`` (or ``embed_invoke_url``) is provided (e.g.
        ``"http://embedding:8000/v1"``), a remote NIM endpoint is used for
        embedding instead of the local HF model.
        """
        resolved = _coerce_params(params, EmbedParams, kwargs)
        if any((resolved.embedding_endpoint, resolved.embed_invoke_url)) and not resolved.api_key:
            resolved = resolved.model_copy(update={"api_key": resolve_remote_api_key()})
        embed_modality = resolved.embed_modality
        embed_granularity = resolved.embed_granularity

        content_columns = (_CONTENT_COLUMNS + ("images",)) if self._caption_enabled else _CONTENT_COLUMNS

        if embed_granularity == "page":
            # Page-level: one row per page with concatenated text and full page image.
            self._tasks.append(
                (
                    collapse_content_to_page_rows,
                    {"modality": embed_modality, "content_columns": content_columns},
                )
            )
        else:
            # Element-level (default): one row per table/chart/infographic.
            text_elements_modality = resolved.text_elements_modality or embed_modality
            structured_elements_modality = resolved.structured_elements_modality or embed_modality
            self._tasks.append(
                (
                    explode_content_to_rows,
                    {
                        "modality": embed_modality,
                        "text_elements_modality": text_elements_modality,
                        "structured_elements_modality": structured_elements_modality,
                        "content_columns": content_columns,
                    },
                )
            )

        from nemo_retriever.params.utils import build_embed_kwargs

        embed_kwargs = build_embed_kwargs(resolved)

        # Ensure embed_modality is forwarded to the embedding function.
        embed_kwargs["embed_modality"] = embed_modality

        # If a remote NIM endpoint is configured, skip local model creation.
        endpoint = (embed_kwargs.get("embedding_endpoint") or embed_kwargs.get("embed_invoke_url") or "").strip()
        if endpoint:
            embed_kwargs.setdefault("input_type", "passage")
            self._tasks.append((embed_text_main_text_embed, embed_kwargs))
            return self

        # Local HF embedder path.
        device = embed_kwargs.pop("device", None)
        hf_cache_dir = embed_kwargs.pop("hf_cache_dir", None)
        normalize = bool(embed_kwargs.pop("normalize", True))
        max_length = int(embed_kwargs.pop("max_length", 8192))
        model_name_raw = embed_kwargs.pop("model_name", None)

        from nemo_retriever.model import create_local_embedder

        embed_kwargs.setdefault("input_type", "passage")
        embed_kwargs["model"] = create_local_embedder(
            model_name_raw,
            device=str(device) if device is not None else None,
            hf_cache_dir=str(hf_cache_dir) if hf_cache_dir is not None else None,
            normalize=normalize,
            max_length=max_length,
        )
        self._tasks.append((embed_text_main_text_embed, embed_kwargs))
        return self

    def save_to_disk(
        self,
        output_directory: Optional[str] = None,
        cleanup: bool = True,
        compression: Optional[str] = "gzip",
    ) -> "InProcessIngestor":
        """
        Persist the current per-document DataFrame to disk as JSON.

        Writes one JSON file per input document under `output_directory`, using a
        UTC timestamp in the filename.
        """
        _ = (cleanup, compression)  # reserved for future use (parity with interface)
        if output_directory is None:
            raise ValueError("output_directory is required for inprocess save_to_disk()")

        self._tasks.append((save_dataframe_to_disk_json, {"output_directory": str(output_directory)}))
        return self

    def vdb_upload(self, params: VdbUploadParams | None = None, **kwargs: Any) -> "InProcessIngestor":
        """
        Upload the (embedding-enriched) results to a vector DB (LanceDB).

        This is an in-process uploader intended to run after `.embed(...)`.
        It reads embeddings from either:
        - `metadata.embedding` (if present), or
        - the embedding payload column produced by the embed stage (default: `text_embeddings_1b_v2`)

        Configuration is passed via kwargs:
        - **lancedb_uri**: str, default `"lancedb"`
        - **table_name**: str, default `"nv-ingest"`
        - **overwrite**: bool, default True (False appends)
        - **create_index**: bool, default True
        - **index_type**: str, default `"IVF_HNSW_SQ"`
        - **metric**: str, default `"l2"`
        - **num_partitions**: int, default 16
        - **num_sub_vectors**: int, default 256
        - **embedding_column**: str, default `"text_embeddings_1b_v2"`
        - **embedding_key**: str, default `"embedding"`
        - **include_text**: bool, default False
        - **text_column**: str, default `"text"`

        Notes:
        - `purge_results_after_upload` is accepted for API parity but is not used in inprocess mode.
        """
        p = params or VdbUploadParams()
        if kwargs:
            lancedb_kwargs = {k: v for k, v in kwargs.items() if k != "purge_results_after_upload"}
            if lancedb_kwargs:
                p = p.model_copy(update={"lancedb": p.lancedb.model_copy(update=lancedb_kwargs)})
            if "purge_results_after_upload" in kwargs:
                p = p.model_copy(update={"purge_results_after_upload": bool(kwargs["purge_results_after_upload"])})
        _ = p.purge_results_after_upload  # parity with interface; inprocess does not purge by default
        self._tasks.append((upload_embeddings_to_lancedb_inprocess, p.lancedb.model_dump(mode="python")))
        return self

    # Tasks that run once on combined results (after all docs). All others run per-doc.
    _POST_TASKS = (upload_embeddings_to_lancedb_inprocess, save_dataframe_to_disk_json)

    def get_pipeline_tasks(
        self,
    ) -> Tuple[List[Tuple[Callable[..., Any], Dict[str, Any]]], List[Tuple[Callable[..., Any], Dict[str, Any]]]]:
        """
        Return (per_doc_tasks, post_tasks) for use by ingest() or by the online
        serve deployment to run the same pipeline on a single document.
        """
        per_doc_tasks = [(f, k) for f, k in self._tasks if f not in self._POST_TASKS]
        post_tasks = [(f, k) for f, k in self._tasks if f in self._POST_TASKS]
        return per_doc_tasks, post_tasks

    def ingest(self, params: IngestExecuteParams | None = None, **kwargs: Any) -> list[Any]:
        run_params = _coerce_params(params, IngestExecuteParams, kwargs)
        show_progress = run_params.show_progress
        _ = (run_params.return_failures, run_params.save_to_disk, run_params.return_traces)
        parallel = run_params.parallel
        max_workers = run_params.max_workers
        gpu_devices = list(run_params.gpu_devices) if run_params.gpu_devices else None
        page_chunk_size = run_params.page_chunk_size

        _start = time.perf_counter()

        # -- Task classification -------------------------------------------
        from nemo_retriever.caption.caption import caption_images as _caption_images_fn

        _post_task_fns = (upload_embeddings_to_lancedb_inprocess, save_dataframe_to_disk_json)
        _cpu_task_fns = (pdf_extraction,)
        # Caption runs on its own device (--caption-device), not in the GPU pool.
        _own_device_fns = (_caption_images_fn,)

        cpu_tasks = [(f, k) for f, k in self._tasks if f in _cpu_task_fns]
        gpu_tasks = [
            (f, k)
            for f, k in self._tasks
            if f not in _cpu_task_fns and f not in _post_task_fns and f not in _own_device_fns
        ]
        own_device_tasks = [(f, k) for f, k in self._tasks if f in _own_device_fns]
        post_tasks = [(f, k) for f, k in self._tasks if f in _post_task_fns]

        docs = list(self._documents)

        # -- Parallel execution branch ------------------------------------
        if parallel:
            if gpu_devices and len(gpu_devices) >= 1 and gpu_tasks:
                # Pipelined: GPU workers load models while CPU runs,
                # each completed chunk goes to GPU immediately.
                from .gpu_pool import GPUWorkerPool, gpu_tasks_to_descriptors

                descriptors = gpu_tasks_to_descriptors(gpu_tasks)
                errors: list[str] = []

                with GPUWorkerPool(gpu_devices, descriptors) as gpu_pool:
                    # Split all docs into page chunks (main thread, cheap PDF byte splitting)
                    chunks: list[pd.DataFrame] = []
                    chunk_to_doc: list[str] = []
                    doc_chunk_total: dict[str, int] = defaultdict(int)
                    for doc in docs:
                        for chunk_df in _iter_page_chunks(doc, page_chunk_size):
                            chunk_to_doc.append(doc)
                            doc_chunk_total[doc] += 1
                            chunks.append(chunk_df)

                    shard_id = 0
                    progress = None
                    if show_progress and tqdm is not None:
                        progress = tqdm(total=len(docs), desc="Processing files", unit="file")

                    shard_to_doc: dict[int, str] = {}
                    doc_done: dict[str, int] = defaultdict(int)

                    def _check_file_done(doc_path: str) -> None:
                        if doc_done[doc_path] >= doc_chunk_total[doc_path]:
                            if progress is not None:
                                progress.update(1)

                    with ProcessPoolExecutor(
                        max_workers=max_workers,
                        mp_context=multiprocessing.get_context("spawn"),
                    ) as cpu_pool:
                        future_to_idx = {
                            cpu_pool.submit(_process_chunk_cpu, chunk, cpu_tasks): i for i, chunk in enumerate(chunks)
                        }

                        for future in as_completed(future_to_idx):
                            idx = future_to_idx[future]
                            doc = chunk_to_doc[idx]
                            try:
                                result = future.result()
                                if isinstance(result, pd.DataFrame) and not result.empty:
                                    for func, kw in own_device_tasks:
                                        result = func(result, **kw)
                                    shard_to_doc[shard_id] = doc
                                    gpu_pool.submit(shard_id, result)
                                    shard_id += 1
                                else:
                                    doc_done[doc] += 1
                                    _check_file_done(doc)
                            except Exception as e:
                                errors.append(f"chunk {idx}: {type(e).__name__}: {e}")
                                print(f"Warning: failed to process chunk {idx}: {e}")
                                doc_done[doc] += 1
                                _check_file_done(doc)

                    if errors:
                        print(f"Warning: {len(errors)} chunk(s) failed CPU extraction")

                    # All CPU done, collect remaining GPU results
                    def _on_gpu_done(sid: int) -> None:
                        d = shard_to_doc.get(sid, "")
                        if d:
                            doc_done[d] += 1
                            _check_file_done(d)

                    combined = gpu_pool.collect_all(on_shard_done=_on_gpu_done)

                    if progress is not None:
                        progress.close()

                if combined.empty:
                    results: list = []
                    if show_progress:
                        _print_ingest_summary(results, time.perf_counter() - _start)
                    return results

                for func, kwargs in post_tasks:
                    combined = func(combined, **kwargs)

                results = [combined]
                if show_progress:
                    _print_ingest_summary(results, time.perf_counter() - _start)
                return results

            else:
                # Single-GPU or no-GPU path: CPU parallel on chunks, then sequential GPU
                chunks: list[pd.DataFrame] = []
                chunk_to_doc: list[str] = []
                doc_chunk_total: dict[str, int] = defaultdict(int)
                for doc in docs:
                    for chunk_df in _iter_page_chunks(doc, page_chunk_size):
                        chunk_to_doc.append(doc)
                        doc_chunk_total[doc] += 1
                        chunks.append(chunk_df)

                cpu_results: list[pd.DataFrame] = []
                errors: list[str] = []

                progress = None
                if show_progress and tqdm is not None:
                    progress = tqdm(total=len(docs), desc="Processing files", unit="file")

                doc_done: dict[str, int] = defaultdict(int)

                with ProcessPoolExecutor(
                    max_workers=max_workers,
                    mp_context=multiprocessing.get_context("spawn"),
                ) as pool:
                    future_to_idx = {pool.submit(_process_chunk_cpu, c, cpu_tasks): i for i, c in enumerate(chunks)}

                    for future in as_completed(future_to_idx):
                        idx = future_to_idx[future]
                        doc = chunk_to_doc[idx]
                        try:
                            result = future.result()
                            if isinstance(result, pd.DataFrame) and not result.empty:
                                cpu_results.append(result)
                        except Exception as e:
                            errors.append(f"chunk {idx}: {type(e).__name__}: {e}")
                            print(f"Warning: failed to process chunk {idx}: {e}")
                        doc_done[doc] += 1
                        if doc_done[doc] >= doc_chunk_total[doc] and progress is not None:
                            progress.update(1)

                if progress is not None:
                    progress.close()

                if errors:
                    print(f"Warning: {len(errors)} chunk(s) failed CPU extraction")

                if not cpu_results:
                    results = []
                    if show_progress:
                        _print_ingest_summary(results, time.perf_counter() - _start)
                    return results

                combined = pd.concat(cpu_results, ignore_index=True)
                for func, kwargs in own_device_tasks:
                    combined = func(combined, **kwargs)
                for func, kwargs in gpu_tasks:
                    combined = func(combined, **kwargs)

                for func, kwargs in post_tasks:
                    combined = func(combined, **kwargs)

                results = [combined]
                if show_progress:
                    _print_ingest_summary(results, time.perf_counter() - _start)
                return results

        # -- Sequential execution branch (default) ------------------------
        use_multi_gpu_seq = gpu_devices and len(gpu_devices) >= 1 and gpu_tasks

        if use_multi_gpu_seq:
            # Pipelined: GPU workers process earlier chunks while CPU
            # extracts later chunks of the same (or next) document.
            cpu_and_extract = [(f, k) for f, k in self._tasks if f in _cpu_task_fns]

            from .gpu_pool import GPUWorkerPool, gpu_tasks_to_descriptors

            descriptors = gpu_tasks_to_descriptors(gpu_tasks)

            with GPUWorkerPool(gpu_devices, descriptors) as gpu_pool:
                shard_id = 0
                shard_to_doc: dict[int, str] = {}
                doc_chunk_total: dict[str, int] = defaultdict(int)
                doc_done: dict[str, int] = defaultdict(int)

                for doc_path in docs:
                    for chunk_df in _iter_page_chunks(doc_path, page_chunk_size):
                        doc_chunk_total[doc_path] += 1
                        current: Any = chunk_df
                        for func, kwargs in cpu_and_extract:
                            if func is pdf_extraction:
                                current = func(pdf_binary=current, **kwargs)
                            else:
                                current = func(current, **kwargs)
                        if isinstance(current, pd.DataFrame) and not current.empty:
                            for func, kw in own_device_tasks:
                                current = func(current, **kw)
                            shard_to_doc[shard_id] = doc_path
                            gpu_pool.submit(shard_id, current)
                            shard_id += 1
                        else:
                            doc_done[doc_path] += 1

                progress = None
                if show_progress and tqdm is not None:
                    # Docs whose chunks all failed CPU are already done
                    already_done = sum(
                        1 for d in docs if d not in doc_chunk_total or doc_done.get(d, 0) >= doc_chunk_total[d]
                    )
                    progress = tqdm(total=len(docs), desc="Processing files", unit="file", initial=already_done)

                def _on_gpu_done(sid: int) -> None:
                    d = shard_to_doc.get(sid, "")
                    if d:
                        doc_done[d] += 1
                        if doc_done[d] >= doc_chunk_total[d] and progress is not None:
                            progress.update(1)

                combined = gpu_pool.collect_all(on_shard_done=_on_gpu_done)

                if progress is not None:
                    progress.close()

            if combined.empty:
                results = []
                if show_progress:
                    _print_ingest_summary(results, time.perf_counter() - _start)
                return results

            for func, kwargs in post_tasks:
                combined = func(combined, **kwargs)

            results = [combined]
            if show_progress:
                _print_ingest_summary(results, time.perf_counter() - _start)
            return results

        # Fully sequential: per-document CPU + GPU + post tasks
        per_doc_tasks = [(f, k) for f, k in self._tasks if f not in _post_task_fns]

        results: list[Any] = []
        doc_iter = docs
        if show_progress and tqdm is not None:
            doc_iter = tqdm(docs, desc="Processing documents", unit="doc")

        if self._pipeline_type == "pdf":

            def _loader(p: str) -> pd.DataFrame:
                """
                Load a document as a per-page DataFrame. For .pdf use pdf_path_to_pages_df.
                For .docx/.pptx convert to PDF via LibreOffice then split (same schema).
                """
                abs_path = os.path.abspath(p)
                ext = os.path.splitext(abs_path)[1].lower()
                if ext in SUPPORTED_EXTENSIONS and ext != ".pdf":
                    try:
                        with open(abs_path, "rb") as f:
                            file_bytes = f.read()
                        pdf_bytes = convert_to_pdf_bytes(file_bytes, ext)
                        pages = _split_pdf_to_single_page_bytes(pdf_bytes)
                        out_rows = [{"bytes": b, "path": abs_path, "page_number": i + 1} for i, b in enumerate(pages)]
                        return pd.DataFrame(out_rows)
                    except BaseException as e:
                        return pd.DataFrame([{"bytes": b"", "path": abs_path, "page_number": 0, "error": str(e)}])
                return pdf_path_to_pages_df(p)

        elif self._pipeline_type == "html":

            def _loader(p: str) -> pd.DataFrame:
                return html_file_to_chunks_df(p, params=HtmlChunkParams(**self._extract_html_kwargs))

        elif self._pipeline_type == "image":

            def _loader(p: str) -> pd.DataFrame:
                from nemo_retriever.image.load import image_file_to_pages_df

                return image_file_to_pages_df(p)

        elif self._pipeline_type == "audio":

            def _loader(p: str) -> pd.DataFrame:
                from nemo_retriever.audio.chunk_actor import audio_path_to_chunks_df

                return audio_path_to_chunks_df(p, params=AudioChunkParams(**self._extract_audio_chunk_kwargs))

        else:

            def _loader(p: str) -> pd.DataFrame:
                return txt_file_to_chunks_df(p, params=TextChunkParams(**self._extract_txt_kwargs))

        for doc_path in doc_iter:
            initial_df = _loader(doc_path)
            current, _ = run_pipeline_tasks_on_df(initial_df, per_doc_tasks, None)
            results.append(current)

        # Run upload/save once on combined results so overwrite=True keeps full corpus.
        if results and all(isinstance(r, pd.DataFrame) for r in results):
            combined = pd.concat(results, ignore_index=True)
            for func, kwargs in post_tasks:
                combined = func(combined, **kwargs)

        if show_progress:
            _print_ingest_summary(results, time.perf_counter() - _start)
        return results
