"""P4 — frozen MTEB regression suite over any backend (Codex 10/11 compliant).

Written against mteb 2.x (2.18.12 pinned): mteb 1.x cannot import under
transformers 5.x (removed AutoModelForVision2Seq), and vLLM requires
transformers 5.x, so both sides of every comparison pin the same mteb 2.x.

Two claims, kept separate:
  --suite regression : frozen, cost-bounded subset for BF16-vs-quant DELTAS —
                       English retrieval + long-doc + held-out MIRACL languages
                       via the HardNegatives variants (full MIRACL corpora are
                       131k/518k docs; the hard-negatives variants cap the
                       corpus per query — cost decision 2026-08-03, disclosed
                       in the model cards).
  --suite parity     : tasks with published upstream BF16 numbers, to validate
                       the harness itself before any delta is trusted.

Every text goes through fixtures.adapter prefixes via prompt_type; results are
cached per-task under runs/mteb/<run-id>/ (mteb ResultCache, only-missing
strategy → resumable) together with a pin manifest.

Usage (arcade, vLLM backend):
  python eval/run_mteb.py --backend vllm --model <path> --suite regression --run-id nvfp4-r1
Usage (HF job / GPU, transformers backend):
  python eval/run_mteb.py --backend transformers --model nvidia/Nemotron-3-Embed-8B-BF16 \
      --revision 8ca3ff38... --device cuda --suite regression --run-id bf16-r1
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
from fixtures.adapter import PREFIXES  # noqa: E402

# Frozen suites. "Task.lang" shorthand → ISO639-3 language filter (mteb 2.x
# takes ISO639-3 in get_tasks and maps to task subsets: swa→sw, tel→te).
SUITES = {
    "regression": [
        "NFCorpus",
        "SciFact",
        "FiQA2018",
        "ArguAna",
        "TRECCOVID",
        "MIRACLRetrievalHardNegatives.swa",  # held-out language (never calibrated)
        "MIRACLRetrievalHardNegatives.tel",  # held-out language (never calibrated)
        "LEMBNarrativeQARetrieval",          # long-document
    ],
    "parity": [
        "NFCorpus",
        "SciFact",
        "TRECCOVID",
    ],
}


def make_adapter_encoder(embed_fn, model_name: str, revision: str | None):
    """mteb 2.x AbsEncoder enforcing the canonical query:/passage: prefixes."""
    import numpy as np
    from mteb.models.abs_encoder import AbsEncoder
    from mteb.models.model_meta import ModelMeta
    from mteb.types import PromptType

    class AdapterEncoder(AbsEncoder):
        mteb_model_meta = ModelMeta.create_empty(
            overwrites=dict(name=model_name, revision=revision or "local")
        )

        def __init__(self, batch_size: int = 16):
            self.batch_size = batch_size

        def encode(self, inputs, *, task_metadata, hf_split, hf_subset,
                   prompt_type=None, **kwargs):
            prefix = (
                PREFIXES["query"]
                if prompt_type == PromptType.query
                else PREFIXES["passage"]
            )
            texts = [prefix + t for batch in inputs for t in batch["text"]]
            # Length-sorted batching: backends that pad to the longest batch
            # member (MLX, transformers) waste most compute on padding for
            # mixed-length corpora. Sort by length, encode, restore order.
            order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
            chunks = []
            for i in range(0, len(order), self.batch_size):
                chunks.append(embed_fn([texts[j] for j in order[i : i + self.batch_size]]))
            emb_sorted = np.concatenate(chunks, axis=0)
            out = np.empty_like(emb_sorted)
            out[order] = emb_sorted
            return out

    return AdapterEncoder


def make_backend(args):
    if args.backend == "vllm":
        import numpy as np
        from vllm import LLM
        from vllm.config import PoolerConfig

        from transformers import AutoTokenizer
        from vllm.inputs import TokensPrompt

        llm = LLM(model=args.model, revision=args.revision, runner="pooling",
                  trust_remote_code=True,
                  enforce_eager=True, max_model_len=args.max_len,
                  gpu_memory_utilization=0.85,
                  pooler_config=PoolerConfig(seq_pooling_type="MEAN"))
        tok = AutoTokenizer.from_pretrained(args.model, revision=args.revision)

        def embed_fn(texts):
            # Truncate at the tokenizer level (matches the transformers
            # reference, which uses truncation=True/max_length=max_len) and
            # feed token IDs — raw text over max_model_len is rejected by
            # vLLM (seen on TRECCOVID: 8193-token doc).
            enc = tok(list(texts), truncation=True, max_length=args.max_len)
            outs = llm.embed(
                [TokensPrompt(prompt_token_ids=ids) for ids in enc["input_ids"]]
            )
            vecs = np.asarray([o.outputs.embedding for o in outs], dtype=np.float32)
            vecs /= np.linalg.norm(vecs, axis=1, keepdims=True)
            return vecs

        return embed_fn

    if args.backend == "transformers":
        from fixtures.pooling_fixtures import transformers_backend

        return transformers_backend(args.model, args.revision, args.device)

    if args.backend == "mlx":
        import importlib.util

        mlx_dir = Path(args.model).expanduser()
        spec = importlib.util.spec_from_file_location(
            "nemotron3_embed_mlx", mlx_dir / "nemotron3_embed_mlx.py"
        )
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        model, tok = mod.load(str(mlx_dir))

        # Batch-level heartbeat: MLX compute runs on the Metal GPU, so the
        # process looks idle to CPU-based liveness checks — without these
        # lines a multi-hour encode is indistinguishable from a hang.
        import time as _time

        state = {"n": 0, "t0": _time.monotonic(), "last": 0.0}

        def embed_fn(texts):
            out = mod.encode(model, tok, list(texts), input_type=None)
            state["n"] += len(texts)
            now = _time.monotonic()
            if now - state["last"] >= 30:
                state["last"] = now
                rate = state["n"] / max(now - state["t0"], 1e-9)
                print(f"[mlx] encoded {state['n']} texts "
                      f"({rate:.2f}/s cumulative)", flush=True)
            return out

        return embed_fn

    raise ValueError(f"unknown backend {args.backend}")


def resolve_tasks(task_names):
    import mteb

    tasks = []
    for name in task_names:
        if "." in name:
            base, lang = name.split(".", 1)
            got = mteb.get_tasks(tasks=[base], languages=[lang])
        else:
            got = mteb.get_tasks(tasks=[name])
        if not got:
            raise SystemExit(f"task resolution failed for {name!r}")
        tasks.extend(got)
    return tasks


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--backend", required=True, choices=["vllm", "transformers", "mlx"])
    ap.add_argument("--model", required=True)
    ap.add_argument("--revision", default=None)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--suite", default="regression", choices=list(SUITES))
    ap.add_argument("--tasks", nargs="*", help="override suite task list")
    ap.add_argument("--run-id", required=True)
    ap.add_argument("--batch-size", type=int, default=16)
    ap.add_argument("--max-len", type=int, default=8192)  # = transformers ref truncation
    args = ap.parse_args()

    import mteb
    from mteb.cache import ResultCache

    out_dir = ROOT / "runs" / "mteb" / args.run_id
    out_dir.mkdir(parents=True, exist_ok=True)

    task_names = args.tasks or SUITES[args.suite]
    tasks = resolve_tasks(task_names)

    encoder_cls = make_adapter_encoder(
        make_backend(args), Path(args.model).name or args.model, args.revision
    )
    model = encoder_cls(batch_size=args.batch_size)

    pins = {
        "run_id": args.run_id,
        "backend": args.backend,
        "model": args.model,
        "revision": args.revision,
        "suite": args.suite,
        "tasks": task_names,
        "resolved": [t.metadata.name for t in tasks],
        "mteb_version": mteb.__version__,
        "max_len": args.max_len,
        "tolerances": json.loads((ROOT / "fixtures" / "tolerances.json").read_text()),
    }
    (out_dir / "pins.json").write_text(json.dumps(pins, indent=2))

    # Per-task loop (Codex P4 review): a timeout/OOM must not lose completed
    # tasks, and failed tasks must fail the job explicitly. After each task the
    # cache is mirrored into artifacts_out and, when NQ_UPLOAD_REPO/NQ_JOB_NAME
    # are set (HF Jobs), uploaded immediately — the launcher's final upload is
    # then a no-op safety net, not the only persistence point.
    import os
    import shutil
    import traceback

    upload_repo = os.environ.get("NQ_UPLOAD_REPO")
    job_name = os.environ.get("NQ_JOB_NAME")

    def checkpoint():
        art = ROOT / "artifacts_out"
        if art.is_dir():
            shutil.copytree(out_dir, art / "mteb" / args.run_id, dirs_exist_ok=True)
            if upload_repo and job_name:
                try:
                    from huggingface_hub import upload_folder

                    upload_folder(
                        repo_id=upload_repo, repo_type="dataset",
                        folder_path=str(art), path_in_repo=job_name,
                        commit_message=f"{args.run_id} incremental checkpoint",
                    )
                except Exception as e:  # retried at next checkpoint
                    print(f"[checkpoint] upload failed (non-fatal now): {e}", flush=True)

    summary, failed = {}, []
    for task in tasks:
        name = task.metadata.name
        print(f"[run_mteb] starting task {name}", flush=True)
        try:
            result = mteb.evaluate(
                model,
                tasks=[task],
                cache=ResultCache(str(out_dir)),      # per-task cache → resumable
                overwrite_strategy="only-missing",
                encode_kwargs={"batch_size": args.batch_size},
                raise_error=True,
                show_progress_bar=True,
            )
            trs = getattr(result, "task_results", []) or []
            if not trs:
                raise RuntimeError(f"no task_results returned for {name}")
            for tr in trs:
                try:
                    summary[tr.task_name] = {"main_score": tr.get_score()}
                except Exception:
                    summary[tr.task_name] = {"scores": getattr(tr, "scores", None)}
            print(f"[run_mteb] finished {name}: {summary.get(name)}", flush=True)
        except Exception:
            traceback.print_exc()
            failed.append(name)
            print(f"[run_mteb] TASK FAILED {name}", flush=True)
        manifest = {
            "run_id": args.run_id,
            "completed": sorted(summary),
            "failed": failed,
            "pending": [t.metadata.name for t in tasks
                        if t.metadata.name not in summary
                        and t.metadata.name not in failed],
        }
        (out_dir / "summary.json").write_text(json.dumps(summary, indent=2, default=str))
        (out_dir / "task_manifest.json").write_text(json.dumps(manifest, indent=2))
        checkpoint()

    print(json.dumps({"run_id": args.run_id, "summary": summary, "failed": failed},
                     indent=2, default=str))
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()