"""
ResuMinder embedding service — Nemotron-3-Embed-1B @ 1024d on ZeroGPU.

Public Space, secret-gated: the gate check runs BEFORE the @spaces.GPU
function, so unauthorised calls are rejected without spending one GPU-second
of quota. The shared secret lives in the Space's Secrets (never in this
repo); callers present it per request.

Contract (kept deliberately dumb — the caller owns text composition):
    embed(payload_json, secret) -> JSON
      payload: {"texts": [...], "is_query": false}
      returns: {"dims": 1024, "vectors": [[...f32 x 1024], ...]}

Vectors are sliced to the first 1024 dims of the model's native 2048 and
re-L2-normalised — the exact configuration that tied the 8B's ceiling on the
production-corpus bake-off (dup@1 0.824 vs 0.836 ± 0.014) while halving
vector storage. Prefixes are the model's own: 'query: ' / 'passage: '.
"""
import json
import os

import gradio as gr
import numpy as np
import spaces
import torch
from loguru import logger
from sentence_transformers import SentenceTransformer

MODEL_ID = "nvidia/Nemotron-3-Embed-1B-BF16"
DIM = 1024
MAX_TOKENS = 2048
MAX_TEXTS = 512          # one call embeds at most this many texts
API_KEY = os.environ.get("API_KEY", "")

model = SentenceTransformer(
    MODEL_ID,
    trust_remote_code=True,
    model_kwargs={"torch_dtype": torch.bfloat16},
)
model.max_seq_length = MAX_TOKENS
# ZeroGPU pattern: a module-level .to('cuda') is intercepted by the spaces
# runtime and materialised when a GPU slice attaches to the decorated call.
try:
    model.to("cuda")
except Exception:
    pass


@spaces.GPU(duration=110)
def _embed_gpu(texts, is_query):
    # Timed INSIDE the decorator: this duration approximates what ZeroGPU
    # bills, so callers can meter real quota spend instead of wall clock
    # (which includes network/queue and runs ~5x hot from the prod side).
    import time
    t0 = time.monotonic()
    prefix = "query: " if is_query else "passage: "
    v = model.encode(
        [prefix + t for t in texts],
        batch_size=32,
        normalize_embeddings=False,
        convert_to_numpy=True,
        show_progress_bar=False,
    ).astype(np.float32)
    v = v[:, :DIM]
    v /= np.linalg.norm(v, axis=1, keepdims=True) + 1e-9
    return v, round(time.monotonic() - t0, 2)


def embed(payload_json: str, secret: str) -> str:
    # Reject BEFORE the GPU decorator: a wrong secret must cost zero quota.
    if not API_KEY or secret != API_KEY:
        return json.dumps({"error": "unauthorised"})
    try:
        payload = json.loads(payload_json)
        logger.debug(" payload: {}", payload)
        texts = payload["texts"]
        logger.debug(" texts: {}", texts)
        is_query = bool(payload.get("is_query", False))
    except Exception as exc:
        return json.dumps({"error": f"bad payload: {exc}"})
    if not isinstance(texts, list) or not texts:
        return json.dumps({"error": "texts must be a non-empty list"})
    if len(texts) > MAX_TEXTS:
        return json.dumps({"error": f"max {MAX_TEXTS} texts per call"})
    texts = [str(t)[:20000] for t in texts]
    v, gpu_s = _embed_gpu(texts, is_query)
    return json.dumps({"dims": DIM, "vectors": v.tolist(), "gpu_s": gpu_s})


demo = gr.Interface(
    fn=embed,
    inputs=[
        gr.Textbox(label="""payload_json (e.g.: {"texts": ["a", "b"]})"""), 
        gr.Textbox(label="API_KEY (hp 6)"),
    ],
    outputs=gr.Textbox(label="result"),
    title="embedding service",
    description=(
        "Private-use embedding endpoint (API_KEY). "
        "Not a public demo — calls without API_KEY are rejected."
    ),
)
demo.queue(max_size=64).launch()