"""ZeroGPU embed + rerank Space — free H200 for the RAG benchmark.

Two API endpoints (call with gradio_client):
- `/embed`  (texts, prompt) -> L2-normalised vectors
- `/rerank` (groups)        -> per-group cross-encoder scores, where groups = [[query, [passages]], ...]
                              MANY queries per call, so the ~per-call overhead is amortised.

Both models load to CPU once at startup, then `.to('cuda')` per call — no re-download, pure GPU
time. `@spaces.GPU(duration=...)` ≈ actual time to keep the 25-min/day PRO quota tight.
"""
import os

import gradio as gr
import numpy as np
import spaces
from sentence_transformers import CrossEncoder, SentenceTransformer

EMB_ID = os.environ.get("EMBED_MODEL_ID", "google/embeddinggemma-300m")  # gated → HF_TOKEN secret
RR_ID = os.environ.get("RERANK_MODEL_ID", "BAAI/bge-reranker-v2-m3")     # open
EMB = SentenceTransformer(EMB_ID, device="cpu", token=os.environ.get("HF_TOKEN"))
EMB.max_seq_length = 512
RR = CrossEncoder(RR_ID, device="cpu", max_length=512)
_LOG = {"emb": 0, "rr": 0}  # cumulative counters surfaced in the Space container logs
print(f"[startup] models loaded: embed={EMB_ID}, rerank={RR_ID}", flush=True)


def _emb_dur(texts, prompt=""):
    return min(120, 15 + len(texts) // 50)  # generous so the GPU task isn't aborted


def _rr_dur(groups):
    n = sum(len(g[1]) for g in (groups or []))
    return min(120, 15 + n // 200)  # cross-encoder is heavier than embed; generous so the task isn't aborted


@spaces.GPU(duration=_emb_dur)
def embed(texts, prompt=""):
    if not texts:
        return []
    _LOG["emb"] += len(texts)
    EMB.to("cuda")
    print(f"[embed] +{len(texts)} texts on {EMB.device} → {_LOG['emb']} total", flush=True)
    v = EMB.encode([(prompt or "") + (t or "")[:6000] for t in texts],
                   normalize_embeddings=True, batch_size=256, convert_to_numpy=True, device="cuda")
    return np.asarray(v, dtype=np.float32).tolist()


@spaces.GPU(duration=_rr_dur)
def rerank(groups):
    """groups: list of [query, [passages]]. Returns list of [scores] (one list per group)."""
    if not groups:
        return []
    _LOG["rr"] += len(groups)
    RR.model.to("cuda")
    print(f"[rerank] +{len(groups)} queries / {sum(len(g[1]) for g in groups)} pairs on {RR.model.device} → {_LOG['rr']} total", flush=True)
    pairs, spans = [], [0]
    for q, ps in groups:
        pairs += [(q, (p or "")[:1800]) for p in ps]
        spans.append(len(pairs))
    scores = RR.predict(pairs, batch_size=128, convert_to_numpy=True) if pairs else np.array([])
    return [list(map(float, scores[spans[i]:spans[i + 1]])) for i in range(len(groups))]


with gr.Blocks() as demo:
    gr.Markdown("ZeroGPU embed + rerank for the RAG benchmark (call /embed and /rerank via gradio_client).")
    ein, epfx, eout = gr.JSON(label="texts"), gr.Textbox(label="prompt", value=""), gr.JSON(label="vectors")
    gr.Button("embed").click(embed, [ein, epfx], eout, api_name="embed")
    rin, rout = gr.JSON(label="groups [[query,[passages]],...]"), gr.JSON(label="scores")
    gr.Button("rerank").click(rerank, rin, rout, api_name="rerank")

if __name__ == "__main__":
    demo.queue(max_size=512, default_concurrency_limit=1).launch()