"""Build the FAISS index on a Modal GPU (offline, dev-only).

Why this exists: the embedding model (Nemotron, 1B) is too heavy for a local
CPU. This script borrows a Modal GPU for the **embedding step only**, then
assembles and persists the FAISS index **locally** — so the live Space still
just *loads* the committed index and never embeds the corpus at runtime.

Split (decided): ingestion (reads ``data/processed/``) and FAISS assembly +
persistence run locally; only the GPU embedding runs on Modal. The remote
function reuses the project's own ``Embedder`` so the vectors are identical to
what the Space's query path produces — same model, revision, prefixes, and
L2-normalization.

The embedding model is ungated, so **no HF token is needed** for this job.

Usage::

    uv sync --group dev                       # installs modal locally
    uv run modal token new                     # one-time auth (opens browser)
    uv run modal run scripts/build_index_modal.py

Output: ``data/processed/index/{index.faiss, chunks.jsonl}``. Commit both
(``index.faiss`` is Git-LFS-tracked) and push to the Space.
"""

from __future__ import annotations

import sys
from pathlib import Path

import modal

# Make the project package importable in the *local* entrypoint below.
REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT))

app = modal.App("pay-equity-index")

# Mirror the Space's embedding deps. Unpinned to match pyproject — this is a
# one-shot offline build, not the served runtime. ``add_local_python_source``
# mounts the project's ``src`` package so the remote fn can import ``Embedder``.
image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch",
        "transformers",
        "sentence-transformers",
        "sentencepiece",
        "accelerate",
        "einops",
        "numpy",
        "pyyaml",
    )
    .add_local_python_source("src")
)


@app.function(gpu="A10G", image=image, timeout=900)
def embed_passages_remote(
    texts: list[str],
    repo: str,
    revision: str | None,
    loader: str,
    query_prefix: str,
    passage_prefix: str,
    normalize: bool,
):
    """Embed all passages on a Modal GPU via the project's ``Embedder``.

    Returns an ``(N, D)`` float32 numpy array (~16 MB for the current corpus),
    serialized back to the local entrypoint.
    """
    from src.backend.indexing.embedder import Embedder

    embedder = Embedder(
        repo=repo,
        revision=revision,
        loader=loader,
        device="cuda",
        query_prefix=query_prefix,
        passage_prefix=passage_prefix,
        normalize=normalize,
    )
    print(f"embedding {len(texts)} passages with {repo} (loader={loader}) on GPU …")
    return embedder.embed_passages(texts)


@app.local_entrypoint()
def main() -> None:
    # Ingestion + persistence are local (no GPU): read data/processed/, build
    # chunks, then write the index next to them after the remote embed returns.
    from src.backend.config import load_models_config
    from src.backend.indexing.ingest import (
        build_corpus,
        load_directive_chunks,
        load_ida_chunks,
    )
    from src.backend.indexing.store import DEFAULT_INDEX_DIR, VectorStore

    cfg = load_models_config().get("embeddings", {}) or {}
    if not cfg.get("repo"):
        raise SystemExit("Set embeddings.repo in configs/models.yaml first.")

    n_dir = len(load_directive_chunks())
    n_ida = len(load_ida_chunks())
    chunks = build_corpus()
    print(f"ingested: directive={n_dir} ida={n_ida} total={len(chunks)}")

    vectors = embed_passages_remote.remote(
        [c.text for c in chunks],
        repo=cfg["repo"],
        revision=cfg.get("revision"),
        loader=cfg.get("loader", "sentence-transformers"),
        query_prefix=cfg.get("query_prefix", ""),
        passage_prefix=cfg.get("passage_prefix", ""),
        normalize=cfg.get("normalize", True),
    )
    dim = vectors.shape[1] if vectors.size else 0

    store = VectorStore()
    store.add(chunks, vectors)
    store.persist()
    print(f"persisted {len(store)} vectors (dim={dim}) → {DEFAULT_INDEX_DIR}")