"""Shared loader for the sae_embed nanobind module (SCRFD + ArcFace).

sae_embed.FaceEmbedder loads both ONNX sessions once and exposes an
embed(path) -> FaceResult method, avoiding the per-process model reload cost
of spawning the embed_faces CLI binary for every image.

resolve_arcface() exposes the same default-resolution logic load_embedder uses,
so a caller can stamp the gallery it is about to write with the model that
actually produced its embeddings (GR-004) — the resolved path, not the CLI
argument, which is often None.
"""

import sys
import os
from pathlib import Path

DEFAULT_ARCFACE = "arcface_w600k_r50.onnx"


def resolve_arcface(models_dir: str, arcface: str | None = None) -> str:
    """The ArcFace/LVFace ONNX path load_embedder would use for these arguments.

    TRACES: GR-004 | SR-001

    Single source of truth for "which model is this", so the stamp written into
    a gallery can never drift from the model loaded."""
    return arcface if arcface else str(Path(models_dir) / DEFAULT_ARCFACE)


def load_embedder(build_dir: str, models_dir: str, arcface: str | None = None,
                   conf: float = 0.5, nms: float = 0.4, max_side: int = 500):
    """Import sae_embed from build_dir and construct a FaceEmbedder.

    Exits with a clear error if the module or models are missing — there is
    no subprocess fallback.
    """
    build_path = Path(build_dir).resolve()
    sys.path.insert(0, str(build_path))
    try:
        import sae_embed
    except ImportError as e:
        sys.exit(
            f"sae_embed module not found in {build_path}: {e}\n"
            f"Build it first: cmake --build {build_dir} --target sae_embed"
        )

    models_path = Path(models_dir)
    detector_path = str(models_path / "scrfd_500m_bnkps.onnx")
    arcface_path = resolve_arcface(models_dir, arcface)
    for model, name in [(detector_path, "SCRFD"), (arcface_path, "ArcFace")]:
        if not Path(model).is_file():
            sys.exit(f"{name} model not found: {model}\nRun: bash scripts/download_models.sh")

    # A TRT-backend build cannot load .onnx; it needs pre-built engines from
    # scripts/build_trt_engines.sh.
    #
    # These are passed only on request. The old comment here claimed they were
    # "ignored by ORT" — they are not. The ORT backend treats an engine path as
    # an instruction and raises, which is the right behaviour (silently ignoring
    # a requested engine would be worse), but it meant that merely HAVING a
    # populated trt_cache/ broke every ORT gallery build in the repo, with an
    # error naming a flag the caller never set.
    use_engines = os.environ.get("SAE_USE_TRT_ENGINES", "") not in ("", "0", "false")
    trt = Path(models_path).parent / "trt_cache"
    det_engine = trt / "scrfd.scrfd_500m_bnkps.640.fp16.engine"
    arc_engine = trt / f"arcface.{Path(arcface_path).stem}.b4.fp16.engine"

    return sae_embed.FaceEmbedder(
        detector_path, arcface_path, conf, nms, max_side,
        str(det_engine) if (use_engines and det_engine.is_file()) else "",
        str(arc_engine) if (use_engines and arc_engine.is_file()) else "",
    )