#!/usr/bin/env python3
"""Reproduce papr-embed-code-1 (developed as papr-embed-v2) MTEB scores from the released package with stock mteb.

  python replicate_mteb.py --model ./papr-embed-code-1-4b --tasks coir --out results/code-1-4b
  python replicate_mteb.py --model Qwen/Qwen3-Embedding-4B --prompts task_prompts.json --tasks ArguAna,NFCorpus --out results/base-4b

Loads the model the way mteb loads Qwen3-Embedding (q3e_instruct_loader: sentence-transformers, last-token
pooling, "Instruct: {instruction}\\nQuery:" on queries only, bare documents) with the package's
task_prompts.json. Everything else comes from the package: float32 weights, max_seq_length 8192
(sentence_bert_config.json), tokenizer-default padding. Batch size is mteb's default, 32. A stock Qwen3
baseline is loaded at mteb's pinned revision with the same prompts, fp32 and max_seq_length 8192.
"""
from __future__ import annotations

import argparse, json, time
from pathlib import Path

COIR = ["AppsRetrieval", "COIRCodeSearchNetRetrieval", "CodeFeedbackMT", "CodeFeedbackST", "CodeSearchNetCCRetrieval",
        "CodeTransOceanContest", "CodeTransOceanDL", "CosQA", "StackOverflowQA", "SyntheticText2SQL"]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, help="package directory, or a Qwen/Qwen3-Embedding-* id for the baseline")
    ap.add_argument("--tasks", required=True, help="'coir' or a comma-separated list of mteb task names")
    ap.add_argument("--out", required=True)
    ap.add_argument("--prompts", help="task prompts JSON (default: <package>/task_prompts.json)")
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--efficient-attn", action="store_true",
                    help="repeat K/V instead of SDPA enable_gqa, so fp32 attention uses the memory-efficient kernel "
                         "(needed below ~40 GB of GPU memory; same math, ~1e-7 relative)")
    ap.add_argument("--token-budget", type=int, default=0,
                    help="replace fixed batches with length-sorted batches of <= N padded tokens (small GPUs)")
    ap.add_argument("--endpoint", help="encode through a deployed endpoint instead of locally: "
                                       "https://<dedicated-dns>/v1/projects/<p>/locations/<r>/endpoints/<id>:predict")
    a = ap.parse_args()
    import inspect, mteb, torch
    from mteb.models.model_implementations.qwen3_models import q3e_instruct_loader
    if a.efficient_attn:
        import transformers.integrations.sdpa_attention as sdpa_mod
        sdpa_mod.use_gqa_in_sdpa = lambda *args, **kwargs: False

    out = Path(a.out); out.mkdir(parents=True, exist_ok=True)
    tasks = COIR if a.tasks.lower() == "coir" else [t.strip() for t in a.tasks.split(",") if t.strip()]
    local = Path(a.model).is_dir()
    prompts = json.load(open(a.prompts or Path(a.model) / "task_prompts.json"))
    if local:
        dev = {"device": "cpu"} if a.endpoint else {}
        model = q3e_instruct_loader(a.model, revision=None, prompts_dict=prompts, **dev)
        prov = json.load(open(Path(a.model) / "PROVENANCE.json"))
        base_meta = mteb.get_model_meta(prov["base"]["repo"])
        model.mteb_model_meta = base_meta.model_copy(update={"name": f"papr-ai/{prov['id']}", "revision": prov["trunk_sha256"][:40]})
    else:
        base_meta = mteb.get_model_meta(a.model)
        model = q3e_instruct_loader(a.model, revision=base_meta.revision, prompts_dict=prompts, max_seq_length=8192)
        model.mteb_model_meta = base_meta  # mteb's prediction writer reads the model name from here
    st = model.model
    hf = st[0].auto_model
    if a.token_budget:  # length-sorted batches of <= token_budget padded tokens (same scheme as the server)
        import numpy as np
        orig = st.encode

        def budget_encode(sentences, prompt=None, batch_size=32, **kw):
            if isinstance(sentences, str) or len(sentences) == 0:
                return orig(sentences, prompt=prompt, batch_size=batch_size, **kw)
            texts = [(prompt or "") + s for s in sentences]
            lens = [len(x) for x in st.tokenizer(texts, truncation=True, max_length=st.max_seq_length)["input_ids"]]
            order = sorted(range(len(texts)), key=lambda i: -lens[i])
            out, i = [None] * len(texts), 0
            while i < len(order):
                top, j = lens[order[i]], i + 1
                while j < len(order) and (j - i + 1) * top <= a.token_budget and j - i < 256:
                    j += 1
                idx = order[i:j]
                embs = orig([sentences[k] for k in idx], prompt=prompt, batch_size=len(idx), **kw)
                for k, e in zip(idx, embs):
                    out[k] = e
                i = j
            return torch.stack(out) if isinstance(out[0], torch.Tensor) else np.stack(out)

        st.encode = budget_encode
    if a.endpoint:  # same strings the memory-server client sends: prompt + text for queries, bare documents
        import concurrent.futures as cf, threading, urllib.error, urllib.request
        import google.auth, google.auth.transport.requests
        import numpy as np
        creds, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
        lock = threading.Lock()

        def token() -> str:
            with lock:
                if not creds.valid:
                    creds.refresh(google.auth.transport.requests.Request())
                return creds.token

        def call(chunk):
            body = json.dumps({"instances": [{"inputs": t} for t in chunk]}).encode()
            for attempt in range(6):
                req = urllib.request.Request(a.endpoint, data=body, headers={
                    "Authorization": f"Bearer {token()}", "Content-Type": "application/json"})
                try:
                    with urllib.request.urlopen(req, timeout=900) as r:
                        return json.loads(r.read())["predictions"]
                except urllib.error.HTTPError as e:
                    if e.code in (429, 500, 502, 503, 504) and attempt < 5:
                        time.sleep(2 ** attempt)
                        continue
                    raise

        def endpoint_encode(sentences, prompt=None, batch_size=32, **kw):
            single = isinstance(sentences, str)
            texts = [(prompt or "") + s for s in ([sentences] if single else sentences)]
            chunks = [texts[i:i + 32] for i in range(0, len(texts), 32)]
            with cf.ThreadPoolExecutor(8) as ex:
                arr = np.asarray([v for part in ex.map(call, chunks) for v in part], dtype=np.float32)
            return arr[0] if single else arr

        st.encode = endpoint_encode
    info = {"model": a.model, "dtype": str(next(hf.parameters()).dtype), "max_seq_length": st.max_seq_length,
            "padding_side": st[0].tokenizer.padding_side, "attn": hf.config._attn_implementation, "batch": a.batch,
            "efficient_attn": a.efficient_attn, "token_budget": a.token_budget, "endpoint": a.endpoint,
            "tf32": torch.backends.cuda.matmul.allow_tf32, "torch": torch.__version__, "mteb": mteb.__version__}
    print(f"[load] {info}", flush=True)
    assert info["dtype"] == "torch.float32" and info["max_seq_length"] == 8192, info

    sig = inspect.signature(mteb.evaluate).parameters
    summary = {}
    for task in tasks:
        t0 = time.time()
        kw = {"encode_kwargs": {"batch_size": a.batch}}
        if "cache" in sig:
            RC = getattr(mteb, "ResultCache", None)
            if RC is None:
                from mteb.cache import ResultCache as RC
            kw["cache"] = RC(cache_path=str(out / "_mteb_cache"))
        if "overwrite_strategy" in sig:
            kw["overwrite_strategy"] = "always"
        if "prediction_folder" in sig:  # rankings kept so CosQA/T2SQL can also be scored by text later
            kw["prediction_folder"] = str(out / f"_pred_{task}")
        res = mteb.evaluate(model, [mteb.get_task(task)], **kw)
        tr = list(getattr(res, "task_results", res))[0]
        subsets = {f"{split}/{r.get('hf_subset', 'default')}": r.get("ndcg_at_10", r.get("main_score"))
                   for split, rows in (tr.scores or {}).items() for r in rows}
        vals = [v for v in subsets.values() if v is not None]
        rec = {"task": task, "ndcg_at_10": sum(vals) / len(vals) if vals else None, "main_score": tr.get_score(),
               "subsets": subsets, "seconds": round(time.time() - t0), "instruction": prompts.get(task), **info}
        json.dump(rec, open(out / f"{task}.json", "w"), indent=1)
        summary[task] = rec["ndcg_at_10"]
        print(f"[task] {task}: nDCG@10 {rec['ndcg_at_10'] * 100:.2f} ({rec['seconds']}s)", flush=True)
    json.dump({"info": info, "ndcg_at_10": summary}, open(out / "summary.json", "w"), indent=1)
    if set(COIR) <= set(summary):
        print(f"[summary] CoIR-10 mean {sum(summary[t] for t in COIR) / 10 * 100:.2f}", flush=True)


if __name__ == "__main__":
    main()