"""Stage 4 — build (and optionally push) images with the data pre-loaded.

Staging happens here, on the host, so the Dockerfiles stay simple:

  * settings + mappings are merged into one `<index>.create.json` per index;
    the dump's own `index.knn.derived_source` (a spelling OpenSearch rejects on
    create) is dropped and every k-NN index is created with derived source off
    (keep_derived_source_off).
  * the elasticdump data file is converted to `_bulk` chunks (to_bulk.py),
    since the OpenSearch image has no Node to run elasticdump with.
"""

from __future__ import annotations

import gzip
import json
import os
import shutil
import subprocess
from pathlib import Path
from typing import Dict, List, Optional

from . import hf_addressing, openapi_spec
from .artifact_store import space_artifact_secrets
from .config import (
    BUILD_DIR,
    REPO_ROOT,
    Config,
    SIGLIP2_NVCF_ENDPOINT_DEFAULT,
    SIGLIP2_NVCF_FUNCTION_ID_DEFAULT,
    TRITONSERVER_BASE_IMAGE_DEFAULT,
)
from .enumerate_assets import load_manifest, save_manifest
from .export import NEO4J_ARTIFACTS, OS_ARTIFACTS
from .to_bulk import convert

BAKE_DIR = BUILD_DIR / "bake"
DOCKER_DIR = REPO_ROOT / "docker"


def _run(cmd: List[str], timeout: int = 900, **kwargs) -> subprocess.CompletedProcess:
    print("  $ " + " ".join(cmd[:8]) + (" ..." if len(cmd) > 8 else ""))
    return subprocess.run(cmd, check=True, timeout=timeout, **kwargs)


def _read_json_gz(path: Path) -> Dict:
    """Load a gzipped elasticdump artifact, undoing double JSON encoding.

    elasticdump is inconsistent about this: `--type=settings` writes the body as
    a JSON *string* containing JSON, while `--type=mapping` writes a plain
    object. Upstream loader.sh normalises both with a `try fromjson catch .`
    step; this is the same defence, applied to whichever form we get.
    """
    with gzip.open(path, "rt") as fh:
        doc = json.load(fh)
    for _ in range(3):
        if not isinstance(doc, str):
            break
        try:
            doc = json.loads(doc)
        except json.JSONDecodeError:
            break
    if not isinstance(doc, dict):
        raise SystemExit(f"{path.name}: expected a JSON object, got {type(doc).__name__}")
    return doc


def _unwrap(doc: Dict, key: str) -> Dict:
    """elasticdump writes {"<index>": {"settings": ...}}; loader.sh also accepts
    the bare {"settings": ...} form. Handle both, like loader.sh does."""
    if key in doc:
        return doc[key]
    if len(doc) == 1:
        inner = next(iter(doc.values()))
        if isinstance(inner, dict) and key in inner:
            return inner[key]
    raise SystemExit(f"could not find {key!r} in dump: keys={list(doc)[:5]}")


def build_create_body(settings_path: Path, mapping_path: Path,
                      number_of_shards: Optional[int] = None) -> Dict:
    """`number_of_shards`, when given, replaces the dumped value: it is fixed
    at index creation, so this is the only point a run indexed with another
    count (e.g. upstream's default 20) can be reloaded with fewer shards."""
    settings = _unwrap(_read_json_gz(settings_path), "settings")
    mappings = _unwrap(_read_json_gz(mapping_path), "mappings")

    index_settings = settings.get("index")
    if isinstance(index_settings, dict):
        # Present in the dump, rejected on create; keep_derived_source_off below
        # sets the value. The dump spells it as a literal dotted key mapping to an object:
        #   "knn.derived_source": {"enabled": "false"}
        # but tolerate the nested spelling as well.
        index_settings.pop("knn.derived_source", None)
        index_settings.pop("knn.derived_source.enabled", None)
        if isinstance(index_settings.get("knn"), dict):
            index_settings["knn"].pop("derived_source", None)
        # Cluster-assigned values that cannot be supplied on create.
        for key in ("uuid", "creation_date", "creation_date_string", "provided_name",
                    "version", "resize", "routing", "history"):
            index_settings.pop(key, None)
    if number_of_shards is not None:
        settings.setdefault("index", {})["number_of_shards"] = str(number_of_shards)
    body = {"settings": settings, "mappings": mappings}
    keep_derived_source_off(body)
    return body


def keep_derived_source_off(body: Dict) -> None:
    """Create a k-NN index with `index.knn.derived_source.enabled: false`, as usdsearch's workers do (and its loader.sh
    since usdsearch !307). With it on (the OpenSearch 3.x default), k-NN before 3.10 can give a document with 150 or
    more nested docs another document's nested vectors on read, and a partial update or an export then stores them
    (opensearch-project/k-NN#3623; run 20261012-9d9f42a: one vector on 1,399 of 1,404 documents). Every restore uses
    this, the Space's seed included. RESTORE_KNN_DERIVED_SOURCE=on leaves the 3.x default."""
    if os.environ.get("RESTORE_KNN_DERIVED_SOURCE", "").strip().lower() in ("on", "true", "1"):
        return
    index_settings = body.setdefault("settings", {}).setdefault("index", {})
    if str(index_settings.get("knn", "")).lower() == "true" or '"knn_vector"' in json.dumps(body.get("mappings", {})):
        index_settings["knn.derived_source.enabled"] = "false"
        print("    k-NN derived source off for this index")


def stage_opensearch(cfg: Config, manifest: Dict) -> Path:
    seed_dir = BAKE_DIR / "opensearch" / "seed"
    if seed_dir.exists():
        shutil.rmtree(seed_dir)
    seed_dir.mkdir(parents=True)

    artifacts = manifest.get("opensearch_artifacts") or {}
    if not artifacts:
        raise SystemExit("manifest has no opensearch_artifacts — run `export` first.")

    ordered: List[str] = []
    for index in sorted(artifacts):
        print(f"  staging {index}")
        body = build_create_body(
            OS_ARTIFACTS / f"{index}.settings.json.gz",
            OS_ARTIFACTS / f"{index}.mapping.json.gz",
            number_of_shards=cfg.opensearch_number_of_shards,
        )
        (seed_dir / f"{index}.create.json").write_text(json.dumps(body))
        chunks = convert(
            OS_ARTIFACTS / f"{index}.data.json.gz",
            seed_dir,
            index,
            chunk_docs=cfg.get_int("BULK_CHUNK_DOCS", 500),
        )
        print(f"    {len(chunks)} bulk chunk(s)")
        ordered.append(index)

    (seed_dir / "indexes.txt").write_text("\n".join(ordered) + "\n")

    context = BAKE_DIR / "opensearch"
    for name in ("Dockerfile.opensearch-baked", "opensearch-seed.sh", "check_bulk.py"):
        shutil.copy2(DOCKER_DIR / name, context / name)
    return context


def stage_neo4j(manifest: Dict) -> Path:
    context = BAKE_DIR / "neo4j"
    dumps = context / "dumps"
    if dumps.exists():
        shutil.rmtree(dumps)
    dumps.mkdir(parents=True)

    artifact = manifest.get("neo4j_artifact")
    if not artifact:
        raise SystemExit("manifest has no neo4j_artifact — run `export` first.")

    src = NEO4J_ARTIFACTS / artifact["file"]
    dest = dumps / "neo4j.dump"
    with gzip.open(src, "rb") as fh, dest.open("wb") as out:
        shutil.copyfileobj(fh, out)
    print(f"  staged {dest.name} ({dest.stat().st_size:,} bytes)")

    for name in ("Dockerfile.neo4j-baked", "neo4j-seeded-entrypoint.sh"):
        shutil.copy2(DOCKER_DIR / name, context / name)
    return context


def _image_names(cfg: Config) -> Dict[str, str]:
    registry = cfg.get("REGISTRY", "nvcr.io/m3sujtetvf5w/usdsearch").rstrip("/")
    names = {
        "opensearch": f"{registry}/{cfg.get('OPENSEARCH_IMAGE_NAME', 'opensearch-simready-central')}",
        "neo4j": f"{registry}/{cfg.get('NEO4J_IMAGE_NAME', 'neo4j-simready-central')}",
    }
    # The combined single-container MVP image (docker/Dockerfile.hf-space)
    # always routes siglip2 to NVCF — it has no local Triton — so it's only
    # buildable when the run itself is already NVCF-backed. cfg is the same
    # object bake/verify/push all load from pipeline.env, so this stays
    # consistent across those stages without threading a separate flag.
    if cfg.siglip2_nvcf:
        names["hf_space"] = f"{registry}/{cfg.get('HF_SPACE_IMAGE_NAME', 'hf-space-simready-central')}"
    # The GPU variant (docker/Dockerfile.hf-space-gpu, local siglip2-triton
    # instead of NVCF) is a separate, optional deploy target — independent
    # of SIGLIP2_BACKEND/siglip2_nvcf above, since it serves the same
    # already-indexed corpus regardless of how the indexing stack itself
    # computed embeddings. Off by default (HF_SPACE_BUILD_GPU_VARIANT):
    # existing pipelines shouldn't pay for an extra ~15 GB build/push they
    # didn't ask for.
    if cfg.hf_space_gpu_variant_enabled:
        names["hf_space_gpu"] = f"{registry}/{cfg.get('HF_SPACE_GPU_IMAGE_NAME', 'hf-space-simready-central-gpu')}"
    return names


#: A repo pinned at more revisions than this gets ONE rule instead of one per
#: revision (url_mapping_rules). Above every repo on the submission lists of
#: 2026-10-15 (most: 19), so their rules are unchanged.
COLLAPSE_REVISIONS_ABOVE = 32

#: USDSEARCH_URL_MAPPING_RULES is one environment string in the Space image.
#: Linux refuses any single env or argv string over 128 KiB (MAX_ARG_STRLEN),
#: so past it every process in the container fails to start; stop the bake
#: well before that.
URL_MAPPING_RULES_MAX_BYTES = 100_000


def url_mapping_rules(hf_repos: List[List[str]], asset_counts: Optional[List[List]] = None) -> List[Dict]:
    """deepsearch-api's USDSEARCH_URL_MAPPING_RULES for the HF repos a run
    indexed (manifest["hf_repos"]: host, repo_type, repo_id, revision): each
    repo's hf:// prefix -> its browsable huggingface.co file-viewer prefix.

    A pinned revision maps `hf://.../resolve/<sha>/` to `.../blob/<sha>/`, so
    the links show exactly what was indexed. `.../blob/main/`, the public form
    an unpinned corpus emits, stays accepted inbound as an alias (never
    emitted), so links handed out before the pin still resolve. Unpinned
    (`main`) keeps the plain `hf://.../<repo>/` -> `.../blob/main/` rule.

    A repo pinned at several revisions (submission lists pin each package
    separately) gives the alias to one of them only — inbound matching is
    first-match, so a second copy would never be reached anyway. It goes to
    the revision with the most packages (`asset_counts`, manifest
    ["hf_repo_asset_counts"]: [host, repo_type, repo_id, revision, n]), the
    one most legacy links point into; without counts, the first listed.

    A repo pinned at more than COLLAPSE_REVISIONS_ABOVE revisions (one upload
    commit per package: 657 for PalatialSim on 2026-10-15) gets one rule,
    `hf://.../resolve/` -> `.../blob/`, which carries every revision through
    unchanged, so the rules stop growing with the pins. It has no
    `.../blob/main/` alias: search_utils refuses a public prefix that is a
    prefix of another, and `.../blob/` is one of `.../blob/main/`.
    """
    counts = {tuple(row[:4]): row[4] for row in asset_counts or []}
    alias_revision: Dict[tuple, str] = {}
    pinned: Dict[tuple, int] = {}
    unpinned = set()
    for host, repo_type, repo_id, revision in hf_repos:
        repo = (host, repo_type, repo_id)
        if not hf_addressing.looks_like_sha(revision):
            unpinned.add(repo)
            continue
        pinned[repo] = pinned.get(repo, 0) + 1
        best = alias_revision.get(repo)
        if best is None or counts.get((*repo, revision), 0) > counts.get((*repo, best), 0):
            alias_revision[repo] = revision
    collapsed = {repo for repo, n in pinned.items() if n > COLLAPSE_REVISIONS_ABOVE and repo not in unpinned}
    emitted = set()
    rules = []
    for host, repo_type, repo_id, revision in hf_repos:
        repo = (host, repo_type, repo_id)
        prefix = hf_addressing.REPO_TYPE_PREFIX.get(repo_type, "")
        web = f"https://{host}/" + "/".join(p for p in (prefix, repo_id) if p)
        if repo in collapsed:
            if repo in emitted:
                continue
            emitted.add(repo)
            rules.append({
                "internal": hf_addressing.build(host=host, repo_type=repo_type, repo_id=repo_id,
                                                revision="", is_folder=True) + "resolve/",
                "public": f"{web}/blob/",
            })
            continue
        rule = {
            "internal": hf_addressing.build(host=host, repo_type=repo_type, repo_id=repo_id,
                                            revision=revision, is_folder=True),
            "public": f"{web}/blob/{revision}/",
        }
        if alias_revision.get((host, repo_type, repo_id)) == revision:
            rule["aliases"] = [f"{web}/blob/main/"]
        rules.append(rule)
    return rules


def hf_space_build_args(cfg: Config, manifest: Dict) -> Dict[str, str]:
    """The `docker/Dockerfile.hf-space` ARGs this run needs, shared between
    the local/CI `docker build` (bake.run, as --build-arg) and pipeline/
    hf_space.py's deploy stage (as HF Space "Variables" — Hugging Face
    passes any Space Variable whose name matches a Dockerfile ARG through as
    a build-arg on its own build, same mechanism, different trigger). One
    dict so the two build-arg lists can't silently drift apart.
    """
    storage_backend = manifest.get("storage_backend", "s3")
    args = {
        # No OPENSEARCH_VERSION: the Space Dockerfiles pin their own
        # (HF_SPACE_OPENSEARCH_VERSION) — their jar patches are tied to it.
        "ARTIFACT_S3_URI": cfg.require("ARTIFACT_S3_URI"),
        "ARTIFACT_S3_REGION": cfg.get("ARTIFACT_S3_REGION", "us-east-2"),
        "CORPUS_NAME": manifest["corpus_name"],
        "ES_NAME": manifest["es_name"],
        # Search never makes a live storage-backend call at serve time (the
        # Dockerfile bakes verification off — see its own comment), but the
        # backend *type* still has to match what the corpus was indexed
        # under: it gates a client-side URL/result code path unrelated to
        # verification, and getting it wrong silently returns zero search
        # results rather than erroring. Verified live.
        "STORAGE_BACKEND_TYPE": storage_backend,
        # Gates the Explorer UI's download feature (nginx.conf.template's
        # __DOWNLOAD_ENABLED__) — true only when this image actually runs a
        # live hf-storage-api server for it to call (storage_api backend);
        # see docs/hf-space-deployment.md.
        "DOWNLOAD_ENABLED": "true" if storage_backend == "storage_api" else "false",
        "SIGLIP2_NVCF_ENDPOINT": cfg.get("SIGLIP2_NVCF_ENDPOINT", SIGLIP2_NVCF_ENDPOINT_DEFAULT),
        "SIGLIP2_NVCF_FUNCTION_ID": cfg.get("SIGLIP2_NVCF_FUNCTION_ID", SIGLIP2_NVCF_FUNCTION_ID_DEFAULT),
    }
    # The hf:// -> https:// link rewrite (url_mapping_rules) comes from the
    # repos this run actually indexed, at their pinned revision, so it can't
    # go stale when the corpus is re-pinned. As a Space Variable it is both
    # the build-arg and the runtime env var. Off for an s3:// corpus.
    rules = url_mapping_rules(manifest.get("hf_repos") or [], manifest.get("hf_repo_asset_counts")) if storage_backend == "storage_api" else []
    rules_json = json.dumps(rules, separators=(",", ":"))
    if len(rules_json.encode()) > URL_MAPPING_RULES_MAX_BYTES:
        raise ValueError(
            f"USDSEARCH_URL_MAPPING_RULES would be {len(rules_json.encode())} bytes ({len(rules)} rules), over "
            f"{URL_MAPPING_RULES_MAX_BYTES}: Linux refuses an env string past 128 KiB, so no process in the Space "
            f"would start. Lower COLLAPSE_REVISIONS_ABOVE or ask the submitter to pin fewer commits.")
    args["USDSEARCH_URL_MAPPING_MODE"] = "rules" if rules else "off"
    args["USDSEARCH_URL_MAPPING_RULES"] = rules_json
    # S3_STORAGE_* is only ever meaningful for an S3-sourced corpus, and
    # threading a source bucket name through for an HF-sourced one would be
    # actively misleading.
    if storage_backend != "storage_api":
        args["S3_STORAGE_BUCKET_NAME"] = cfg.get("S3_STORAGE_BUCKET_NAME", "")
        args["S3_STORAGE_REGION_NAME"] = cfg.get("S3_STORAGE_REGION_NAME", "us-west-2")
    return args


def hf_space_gpu_build_args(cfg: Config, manifest: Dict) -> Dict[str, str]:
    """The `docker/Dockerfile.hf-space-gpu` ARGs this run needs — same shape
    as hf_space_build_args() above, minus every NVCF-only key (that
    Dockerfile has no SIGLIP2_NVCF_* ARGs at all: SigLIP2 comes from the
    local siglip2-triton process it bakes in, not NVCF), plus the bare
    Triton base image, overridable in case a newer NGC tritonserver tag
    needs picking up without a Dockerfile edit. The SigLIP2 model weights
    themselves are not a build-arg — they're vendored source
    (docker/hf-space/siglip2-model/, see scripts/extract_siglip2_model.sh),
    committed like every other file this Dockerfile COPYs.
    """
    args = hf_space_build_args(cfg, manifest)
    args.pop("SIGLIP2_NVCF_ENDPOINT", None)
    args.pop("SIGLIP2_NVCF_FUNCTION_ID", None)
    args["TRITONSERVER_BASE_IMAGE"] = cfg.get("TRITONSERVER_BASE_IMAGE", TRITONSERVER_BASE_IMAGE_DEFAULT)
    return args


def image_tags(cfg: Config, manifest: Dict) -> Dict[str, str]:
    # <corpus name>-<run_id> — in CI, run_id is itself
    # <branch slug>-<workflow run id>-<short commit sha> (see
    # nvcf-pipeline.yml's "Compute version" step), so the full tag traces
    # a published image back to the exact corpus, branch, CI run and commit
    # that built it. Locally, run_id falls back to enumerate_assets.py's
    # `YYYYMMDD-<corpus sha prefix>`.
    version = f"{manifest['corpus_name']}-{manifest['run_id']}"
    return {name: f"{image}:{version}" for name, image in _image_names(cfg).items()}


def latest_tags(cfg: Config, manifest: Dict) -> Dict[str, str]:
    """`<corpus name>-latest` — updated only when `push` actually publishes."""
    tag = f"{manifest['corpus_name']}-latest"
    return {name: f"{image}:{tag}" for name, image in _image_names(cfg).items()}


def run(cfg: Config, push: bool = False) -> int:
    manifest = load_manifest()
    tags = image_tags(cfg, manifest)
    common = [
        "--build-arg", f"RUN_ID={manifest['run_id']}",
        "--build-arg", f"CORPUS_SHA256={manifest['corpus_sha256']}",
        "--build-arg", f"USDSEARCH_GIT_SHA={manifest.get('usdsearch_git_sha', 'unknown')}",
    ]

    print("Staging OpenSearch seed ...")
    os_context = stage_opensearch(cfg, manifest)
    print(f"\nBuilding {tags['opensearch']}")
    _run([
        "docker", "build",
        "-f", str(os_context / "Dockerfile.opensearch-baked"),
        "-t", tags["opensearch"],
        "--build-arg", f"OPENSEARCH_VERSION={cfg.get('OPENSEARCH_VERSION', '3.8.0')}",
        "--build-arg", f"MAIN_INDEX={cfg.main_index}",
        "--build-arg", f"DOC_COUNT={manifest.get('base_docs_indexed', 0)}",
        *common, str(os_context),
    ])

    if "hf_space" in tags:
        print(f"\nBuilding {tags['hf_space']} (single-container MVP: opensearch+redis+deepsearch-api, "
              f"for Hugging Face Spaces)")
        # Fails fast here, before the Neo4j build, if no NVCF token is in the
        # environment — same fallback chain as the indexing/runtime nvcf
        # overlays (config.py's siglip2_nvcf_env), so a missing token is
        # reported the same way everywhere rather than surfacing later as an
        # opaque docker build failure. No AWS-credential check here: this
        # image no longer bakes the index in at build time (see the
        # Dockerfile's header) — the artifact-bucket keys are only needed at
        # `verify_baked`'s smoke-test `docker run`, and at the real Space's
        # own runtime.
        cfg.siglip2_nvcf_env()
        # Build context is REPO_ROOT here (not DOCKER_DIR, unlike the
        # opensearch/neo4j builds above): this Dockerfile COPYs pipeline/ in
        # directly (its runtime seed step needs it) rather than reusing
        # another already-built image's tag, since HF's own builder must be
        # able to build it standalone with only public base images.
        build_arg_flags = [
            flag
            for key, value in hf_space_build_args(cfg, manifest).items()
            for flag in ("--build-arg", f"{key}={value}")
        ]
        _run([
            "docker", "build",
            "-f", str(DOCKER_DIR / "Dockerfile.hf-space"),
            "-t", tags["hf_space"],
            *build_arg_flags,
            *common, str(REPO_ROOT),
        ])
    else:
        print("\nSIGLIP2_BACKEND=local — skipping the HF Space image (it is NVCF-only by "
              "design; set SIGLIP2_BACKEND=nvcf to build it — see docker/Dockerfile.hf-space)")

    if "hf_space_gpu" in tags:
        print(f"\nBuilding {tags['hf_space_gpu']} (single-container MVP + local siglip2-triton, "
              f"for a GPU-tier Hugging Face Space — see docker/Dockerfile.hf-space-gpu)")
        # No NVCF token needed here (this variant never calls NVCF) — only
        # the same build-arg dict shape as the NVCF image, minus its
        # NVCF-only keys. Build context is REPO_ROOT, same reasoning as the
        # NVCF hf_space build above.
        build_arg_flags = [
            flag
            for key, value in hf_space_gpu_build_args(cfg, manifest).items()
            for flag in ("--build-arg", f"{key}={value}")
        ]
        _run([
            "docker", "build",
            "-f", str(DOCKER_DIR / "Dockerfile.hf-space-gpu"),
            "-t", tags["hf_space_gpu"],
            *build_arg_flags,
            *common, str(REPO_ROOT),
        ])

    print("\nStaging Neo4j dump ...")
    neo_context = stage_neo4j(manifest)
    print(f"\nBuilding {tags['neo4j']}")
    _run([
        "docker", "build",
        "-f", str(neo_context / "Dockerfile.neo4j-baked"),
        "-t", tags["neo4j"],
        "--build-arg", f"NEO4J_VERSION={cfg.get('NEO4J_VERSION', '5.23.0-community')}",
        "--build-arg", f"NODE_COUNT={manifest.get('neo4j_node_count', 0)}",
        "--build-arg", f"RELATIONSHIP_COUNT={manifest.get('neo4j_relationship_count', 0)}",
        *common, str(neo_context),
    ])

    manifest["images"] = tags
    (BUILD_DIR / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")

    print("\nVerifying the baked data from a fresh container ...")
    if verify_baked(cfg, tags, manifest) != 0:
        return 1

    if push:
        # Only reached after `verify` has passed; run_pipeline never pushes
        # before the smoke test.
        for name, tag in tags.items():
            print(f"\nPushing {tag}")
            _run(["docker", "push", tag])

    print("\nbuilt:")
    for name, tag in tags.items():
        print(f"  {name}: {tag}")
    if not push:
        print("\nNot pushed. `run_pipeline.py verify` runs the smoke test; push with")
        print("`run_pipeline.py push` once it passes.")
    return 0


def download_smoke(base_url: str, timeout_s: float = 180.0) -> bool:
    """A storage_api corpus's download route answers for a real asset: search
    one hit, then ask /download/asset for its manifest (sizes only, no bytes) —
    once as the default whole-package download (package-download, through
    hf-storage-api's folder listing: scope "package", the asset itself among
    the files) and once as scope=dependencies (forwarded to deepsearch-api's
    own asset + dependencies bundle). Polls, because hf-storage-api and
    package-download start only once the seed restored the catalog."""
    import time
    import httpx

    def preview(asset_url: str, **params):
        resp = httpx.get(f"{base_url}/download/asset",
                         params={"asset_url": asset_url, "manifest_only": "true", **params}, timeout=60.0)
        return resp.json() if resp.status_code == 200 else None

    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        try:
            hits = httpx.get(f"{base_url}/search", params={"query": "box", "limit": 1}, timeout=15.0).json()
            asset_url = hits[0]["url"] if hits else None
            if asset_url:
                package = preview(asset_url)
                deps = preview(asset_url, scope="dependencies")
                if (package and package.get("scope") == "package" and (package.get("file_count") or 0) > 0
                        and asset_url in {f["url"] for f in package["files"]}
                        and deps and (deps.get("file_count") or 0) > 0):
                    return True
        except (httpx.HTTPError, ValueError, KeyError, IndexError, TypeError, AttributeError):
            pass
        time.sleep(3)
    return False


def public_url_smoke(base_url: str, public_prefixes: List[str], timeout_s: float = 180.0) -> tuple:
    """Every asset URL the Space hands out is a public huggingface.co link
    under this run's URL-mapping rule: v2 search `url`s, v3 hybrid
    `base_key`s and a download manifest's file URLs. A bare `https://` check
    is not enough. The stale hand-set rule emitted
    `https://.../blob/main/resolve/<sha>/...` (a 404), which passes it, so
    each URL must start with one of `public_prefixes`
    (url_mapping_rules()'s `public`s). Under a collapsed repo's `.../blob/`
    the next segment must be a commit sha, or that stale form would pass.
    Returns (ok, detail). Polls like download_smoke(): the seed restores the
    index after /health is up."""
    import re
    import time
    import httpx

    def under(u):
        return any(u.startswith(p) and (not p.endswith("/blob/") or re.match(r"[0-9a-f]{40}/", u[len(p):]))
                   for p in public_prefixes)

    def bad(urls):
        return [u for u in urls if not (isinstance(u, str) and u.startswith("https://") and under(u))]

    detail = "no URLs returned"
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        try:
            v2 = httpx.get(f"{base_url}/search", params={"description": "", "limit": 25, "return_images": "false"},
                           timeout=30.0).json()
            v2 = v2 if isinstance(v2, list) else v2.get("hits", [])
            v3 = httpx.post(f"{base_url}/search_hybrid", json={"description": "", "limit": 25, "return_images": False},
                            timeout=30.0).json().get("hits", [])
            urls = [h.get("url") for h in v2] + [(h.get("source") or {}).get("base_key") for h in v3]
            if urls and urls[0]:
                files = httpx.get(f"{base_url}/download/asset", params={"asset_url": urls[0], "manifest_only": "true"},
                                  timeout=60.0).json().get("files", [])
                urls += [f.get("url") for f in files]
                wrong = bad(urls)
                if not wrong and v2 and v3 and files:
                    return True, f"{len(urls)} URLs, all under {', '.join(public_prefixes)}"
                if wrong:
                    return False, f"{len(wrong)}/{len(urls)} URLs not under {', '.join(public_prefixes)}, e.g. {wrong[0]}"
                detail = f"incomplete response (v2 {len(v2)}, v3 {len(v3)}, download files {len(files)})"
        except (httpx.HTTPError, ValueError, KeyError, IndexError, TypeError, AttributeError) as exc:
            detail = f"{type(exc).__name__}: {exc}"
        time.sleep(3)
    return False, detail


def api_docs_smoke(base_url: str, container: str, timeout_s: float = 180.0) -> tuple:
    """The Space's API docs (/openapi.json, which /docs and /redoc load) match the API it serves:

    - deepsearch-api's and asset-graph-service's own specs, read inside the container, have no route the docs
      neither document nor leave out on purpose, and still have every documented one (openapi_spec.route_problems:
      an API added, renamed or removed upstream needs a decision here first);
    - /openapi.json is pipeline/openapi_spec.py's merge of exactly those specs, not deepsearch-api's own;
    - every documented operation, called through nginx without parameters (GET) or with ``{}`` (POST), reaches its
      app: a JSON answer (a 200, or the app's own 400/422), never the Explorer's HTML, nginx's 404/405 or a 5xx.

    Returns (ok, detail). Polls: the openapi-spec program writes the file once both apps are up."""
    import time
    import httpx

    def internal(url):
        out = subprocess.run(["docker", "exec", container, "curl", "-sf", url],
                             capture_output=True, text=True, timeout=60)
        return json.loads(out.stdout) if out.returncode == 0 else None

    detail = "no spec yet"
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        try:
            search, ags = internal(openapi_spec.SEARCH_URL), internal(openapi_spec.AGS_URL)
            served = httpx.get(f"{base_url}/openapi.json", timeout=30.0).json()
            if search and ags and served.get("info", {}).get("title") == openapi_spec.TITLE:
                break
            detail = "/openapi.json is still deepsearch-api's own spec"
        except (httpx.HTTPError, ValueError, subprocess.SubprocessError) as exc:
            detail = f"{type(exc).__name__}: {exc}"
        time.sleep(3)
    else:
        return False, detail

    problems = openapi_spec.route_problems(search, ags)
    if problems:
        return False, f"{len(problems)} route(s) the docs don't cover: " + "; ".join(problems)
    expected = openapi_spec.merge(search, ags)
    if set(served["paths"]) != set(expected["paths"]):
        return False, f"/openapi.json paths differ from the merge: {sorted(set(served['paths']) ^ set(expected['paths']))}"
    wrong = []
    for path, item in served["paths"].items():
        for method in item:
            if method not in ("get", "post"):
                continue
            try:
                resp = httpx.request(method.upper(), f"{base_url}{path}", json={} if method == "post" else None,
                                     timeout=60.0)
                ctype = resp.headers.get("content-type", "")
                if resp.status_code in (404, 405) or resp.status_code >= 500 or not ctype.startswith("application/json"):
                    wrong.append(f"{method.upper()} {path} -> {resp.status_code} {ctype}")
            except httpx.HTTPError as exc:
                wrong.append(f"{method.upper()} {path} -> {type(exc).__name__}")
    if wrong:
        return False, f"{len(wrong)} documented operation(s) don't reach their app: " + "; ".join(wrong)
    return True, f"{len(served['paths'])} documented paths, each reaching its app"


def verify_baked(cfg: Config, tags: Dict[str, str], manifest: Dict) -> int:
    """Check the data from a fresh `docker run`, not from the build log.

    This is what catches the VOLUME trap: a Dockerfile that writes into a
    declared VOLUME builds green and ships an empty database, and only starting
    the built image reveals it.
    """
    ok = True

    print("  opensearch ...")
    result = subprocess.run(
        ["docker", "run", "--rm", "-d", "-p", "19200:9200", tags["opensearch"]],
        capture_output=True, text=True, check=True, timeout=60,
    )
    container = result.stdout.strip()
    try:
        import time
        import httpx
        expected = manifest.get("base_docs_indexed", 0)
        count = -1
        shards_ready = False
        for _ in range(90):
            try:
                response = httpx.get(f"http://localhost:19200/{cfg.main_index}/_count", timeout=5.0)
                if response.status_code < 400:
                    body = response.json()
                    shards = body.get("_shards", {})
                    shards_ready = shards.get("successful") == shards.get("total")
                    count = int(body["count"])
                    # A fresh container recovers its shards from the baked
                    # segments over several seconds. OpenSearch answers
                    # `_count` as soon as any shard is open, and
                    # `allow_partial_search_results` defaults to true, so a
                    # response with `shards_ready=False` can silently
                    # undercount rather than error — keep polling instead of
                    # trusting the first HTTP 200.
                    if shards_ready and count == expected:
                        break
            except httpx.HTTPError:
                pass
            time.sleep(2)
        print(f"    {cfg.main_index}: {count} docs (expected {expected}, shards ready: {shards_ready})")
        if count != expected or count == 0:
            print("    FAIL: baked OpenSearch image does not hold the expected documents")
            ok = False
    finally:
        subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False, timeout=30)

    print("  neo4j ...")
    result = subprocess.run(
        ["docker", "run", "--rm", "-d", "-p", "17687:7687", tags["neo4j"]],
        capture_output=True, text=True, check=True, timeout=60,
    )
    container = result.stdout.strip()
    try:
        import time
        from neo4j import GraphDatabase
        expected = manifest.get("neo4j_node_count", 0)
        nodes = -1
        # A bounded wall-clock deadline, not a bare iteration count: the
        # driver is given no timeouts below by default, so a single stuck
        # connection attempt (e.g. the server accepting TCP but stalling
        # mid-handshake) can block far longer than the intended ~3 minutes
        # of polling — 90 iterations at an unbounded per-attempt cost was
        # exactly the shape of a run that hung here for 44 minutes and
        # blew straight through this step's own 10-minute CI timeout
        # without ever failing cleanly.
        deadline = time.monotonic() + 180
        while time.monotonic() < deadline:
            try:
                driver = GraphDatabase.driver(
                    "bolt://localhost:17687", auth=None,
                    connection_timeout=5.0, connection_acquisition_timeout=5.0,
                )
                with driver.session() as session:
                    nodes = session.run("MATCH (n) RETURN count(n) AS c").single()["c"]
                driver.close()
                break
            except Exception:  # noqa: BLE001 - the server is simply not up yet
                time.sleep(2)
        print(f"    nodes: {nodes} (expected {expected})")
        if expected > 0 and nodes != expected:
            print("    FAIL: baked Neo4j image does not hold the expected graph")
            ok = False
    finally:
        subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False, timeout=30)

    if "hf_space" in tags:
        print("  hf-space (combined MVP image) ...")
        # The AWS keys here exercise the same runtime restore
        # (pipeline/hf_space_seed.py, via supervisord) a real Space does —
        # this smoke test is what actually proves that path works, not just
        # that the image builds. A longer deadline than the other images'
        # health checks: startup now includes a real S3 download + bulk
        # load, not just a process coming up. RESTORE_RUN_ID pins the restore
        # to this run: latest.json may name another one (`upload --no-latest`).
        token = cfg.siglip2_nvcf_env()["SIGLIP2_TRITON_SERVER_AUTH_TOKEN"]
        aws = space_artifact_secrets()  # the read-only keys the Space itself gets
        result = subprocess.run(
            ["docker", "run", "--rm", "-d", "-p", "17860:7860",
             "-e", f"SIGLIP2_TRITON_SERVER_AUTH_TOKEN={token}",
             *(flag for key, value in aws.items() for flag in ("-e", f"{key}={value}")),
             "-e", f"RESTORE_RUN_ID={manifest['run_id']}",
             tags["hf_space"]],
            capture_output=True, text=True, check=True, timeout=60,
        )
        container = result.stdout.strip()
        try:
            import time
            import httpx
            healthy = False
            deadline = time.monotonic() + 300
            while time.monotonic() < deadline:
                try:
                    if httpx.get("http://localhost:17860/health", timeout=5.0).status_code == 200:
                        healthy = True
                        break
                except httpx.HTTPError:
                    pass
                time.sleep(2)
            print(f"    /health responded: {healthy}")
            if not healthy:
                print("    FAIL: hf-space image never became healthy")
                ok = False

            # /health only proves deepsearch-api/OpenSearch came up — it says
            # nothing about neo4j/asset-graph-service, which have their own
            # offline-load startup path (pipeline/hf_space_seed.py
            # --component=neo4j) that can fail silently as far as /health is
            # concerned, and which routinely takes longer than deepsearch-api's
            # own startup (a multi-hundred-MB neo4j-admin load, run offline,
            # before that program even execs into `neo4j console`) — so this
            # has to poll, not fire once right after /health passes. Exercise
            # a real graph read through nginx, the same way a client actually
            # would: search for a real asset url, then ask
            # asset-graph-service for its prims.
            graph_ok = False
            if healthy:
                graph_deadline = time.monotonic() + 180
                while time.monotonic() < graph_deadline and not graph_ok:
                    try:
                        search_resp = httpx.get(
                            "http://localhost:17860/search", params={"query": "box", "limit": 1}, timeout=10.0,
                        )
                        results = search_resp.json()
                        asset_url = results[0]["url"] if results else None
                        if asset_url:
                            prims_resp = httpx.get(
                                "http://localhost:17860/asset_graph/usd/prims",
                                params={"asset_url": asset_url}, timeout=10.0,
                            )
                            prims = prims_resp.json()
                            graph_ok = prims_resp.status_code == 200 and isinstance(prims, list) and len(prims) > 0
                    except (httpx.HTTPError, ValueError):
                        pass
                    if not graph_ok:
                        time.sleep(3)
                print(f"    asset-graph-service responded with real prims: {graph_ok}")
                if not graph_ok:
                    print("    FAIL: hf-space image's neo4j/asset-graph-service restore did not produce queryable data")
                    ok = False

            # Recorded independently of `ok`/ the other images' checks —
            # deploy-space (pipeline/hf_space.py) gates on this specifically,
            # not on the separate full-stack `verify` stage's smoke_passed,
            # since that stage also exercises asset-graph-service and other
            # components the hf_space image doesn't even include.
            download_ok = True
            if healthy and manifest.get("storage_backend") == "storage_api":
                download_ok = download_smoke("http://localhost:17860")
                print(f"    /download/asset package + dependencies manifests for a real asset: {download_ok}")
                if not download_ok:
                    print("    FAIL: hf-space image's download route (package-download / hf-storage-api) does not answer")
                    ok = False

            # Returned asset URLs are this run's public huggingface.co links
            # (USDSEARCH_URL_MAPPING_RULES), so a stale or missing rule can't ship.
            urls_ok = True
            rules = url_mapping_rules(manifest.get("hf_repos") or [], manifest.get("hf_repo_asset_counts"))
            if healthy and manifest.get("storage_backend") == "storage_api" and rules:
                urls_ok, detail = public_url_smoke("http://localhost:17860", [r["public"] for r in rules])
                print(f"    returned asset URLs are public https:// links: {urls_ok} ({detail})")
                if not urls_ok:
                    print("    FAIL: hf-space image returns asset URLs outside this run's public URL mapping")
                    ok = False

            # /docs and /redoc show the API the Space serves (api_docs_smoke): an API added, renamed or removed
            # in usdsearch fails here until pipeline/openapi_spec.py and nginx.conf.template account for it.
            docs_ok = False
            if healthy:
                docs_ok, detail = api_docs_smoke("http://localhost:17860", container)
                print(f"    /openapi.json documents the served API: {docs_ok} ({detail})")
                if not docs_ok:
                    print("    FAIL: hf-space image's API docs don't match its API (docs/development.md, "
                          "\"The Space's API docs\")")
                    ok = False

            manifest["hf_space_smoke_passed"] = healthy and graph_ok and download_ok and urls_ok and docs_ok
            save_manifest(manifest)
        finally:
            subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False, timeout=30)

    if "hf_space_gpu" in tags:
        print("  hf-space-gpu (local siglip2-triton variant) ...")
        # Unlike every other smoke test in this function, this one needs an
        # actual GPU on the machine running `bake` — `--gpus all` fails fast
        # and loudly if there isn't one, rather than hanging, which is the
        # right failure mode here: this is why this variant is opt-in
        # (HF_SPACE_BUILD_GPU_VARIANT) and why the existing nightly CI
        # workflow (no-GPU-runner, per docs/plans/index-regeneration-
        # pipeline.md §12) doesn't enable it. No NVCF token needed. AWS keys
        # still are — the index restore path is identical to the NVCF
        # variant's.
        aws = space_artifact_secrets()  # the read-only keys the Space itself gets
        result = subprocess.run(
            ["docker", "run", "--rm", "-d", "-p", "17861:7860", "--gpus", "all",
             *(flag for key, value in aws.items() for flag in ("-e", f"{key}={value}")),
             "-e", f"RESTORE_RUN_ID={manifest['run_id']}",
             tags["hf_space_gpu"]],
            capture_output=True, text=True, check=True, timeout=60,
        )
        container = result.stdout.strip()
        try:
            import time
            import httpx
            healthy = False
            # Longer deadline than the NVCF variant's: this one also has to
            # load ~3.6 GB of ONNX weights onto the GPU before deepsearch-api
            # execs uvicorn (see supervisord-gpu.conf) — not measured live
            # yet (see docker/Dockerfile.hf-space-gpu's header), budgeted
            # generously rather than guessed tight.
            deadline = time.monotonic() + 600
            while time.monotonic() < deadline:
                try:
                    if httpx.get("http://localhost:17861/health", timeout=5.0).status_code == 200:
                        healthy = True
                        break
                except httpx.HTTPError:
                    pass
                time.sleep(2)
            print(f"    /health responded: {healthy}")
            if not healthy:
                print("    FAIL: hf-space-gpu image never became healthy")
                ok = False

            embedding_ok = False
            if healthy:
                try:
                    search_resp = httpx.get(
                        "http://localhost:17861/search", params={"query": "box", "limit": 1}, timeout=15.0,
                    )
                    results = search_resp.json()
                    # A non-empty, non-erroring hybrid-search response through
                    # nginx is proof the local siglip2-triton text-embedding
                    # round trip actually worked, not just that the process
                    # answered /health — same "exercise a real call, not just
                    # readiness" reasoning as the NVCF variant's asset-graph
                    # check above.
                    embedding_ok = search_resp.status_code == 200 and isinstance(results, list) and len(results) > 0
                except (httpx.HTTPError, ValueError):
                    pass
                print(f"    local siglip2-triton search round trip: {embedding_ok}")
                if not embedding_ok:
                    print("    FAIL: hf-space-gpu image's local siglip2-triton embedding path did not "
                          "return real results")
                    ok = False

            download_ok = True
            if healthy and manifest.get("storage_backend") == "storage_api":
                download_ok = download_smoke("http://localhost:17861")
                print(f"    /download/asset manifest for a real asset: {download_ok}")
                if not download_ok:
                    print("    FAIL: hf-space-gpu image's download route (hf-storage-api) does not answer")
                    ok = False

            # Returned asset URLs are this run's public huggingface.co links
            # (USDSEARCH_URL_MAPPING_RULES), so a stale or missing rule can't ship.
            urls_ok = True
            rules = url_mapping_rules(manifest.get("hf_repos") or [], manifest.get("hf_repo_asset_counts"))
            if healthy and manifest.get("storage_backend") == "storage_api" and rules:
                urls_ok, detail = public_url_smoke("http://localhost:17861", [r["public"] for r in rules])
                print(f"    returned asset URLs are public https:// links: {urls_ok} ({detail})")
                if not urls_ok:
                    print("    FAIL: hf-space-gpu image returns asset URLs outside this run's public URL mapping")
                    ok = False

            # /docs and /redoc show the API the Space serves (api_docs_smoke): an API added, renamed or removed
            # in usdsearch fails here until pipeline/openapi_spec.py and nginx.conf.template account for it.
            docs_ok = False
            if healthy:
                docs_ok, detail = api_docs_smoke("http://localhost:17861", container)
                print(f"    /openapi.json documents the served API: {docs_ok} ({detail})")
                if not docs_ok:
                    print("    FAIL: hf-space-gpu image's API docs don't match its API (docs/development.md, "
                          "\"The Space's API docs\")")
                    ok = False

            manifest["hf_space_gpu_smoke_passed"] = healthy and embedding_ok and download_ok and urls_ok and docs_ok
            save_manifest(manifest)
        finally:
            subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False, timeout=30)

    return 0 if ok else 1