"""Check vLLM LoRA embeddings for a BERT-family model against HF + PEFT.

Usage:
  python validate_bert_lora.py make-adapter MODEL OUT_DIR
  python validate_bert_lora.py prefix-adapter SRC_DIR OUT_DIR PREFIX
  python validate_bert_lora.py check MODEL ADAPTER_DIR [--hf-adapter DIR]
                                     [--expect-unsupported]
"""

import json
import sys

import torch
import torch.nn.functional as F

TEXTS = [
    "The quick brown fox jumps over the lazy dog.",
    "vLLM serves embedding models with paged attention.",
    "LoRA adds a low-rank update to frozen weights.",
    "Blacksburg is home to Virginia Tech.",
]


def make_adapter(model_name: str, out_dir: str) -> None:
    from peft import LoraConfig, get_peft_model
    from transformers import AutoModel

    torch.manual_seed(0)
    base = AutoModel.from_pretrained(model_name, torch_dtype=torch.float32)
    config = LoraConfig(
        r=8,
        lora_alpha=16,
        # Encoder projections only; HF's unused `pooler.dense` is left out.
        target_modules=r".*encoder.*\.(query|key|value|dense)$",
        init_lora_weights=False,  # random A and B, so the adapter has an effect
    )
    peft_model = get_peft_model(base, config)
    peft_model.save_pretrained(out_dir)
    from safetensors import safe_open

    with safe_open(f"{out_dir}/adapter_model.safetensors", "pt") as f:
        names = sorted(f.keys())
    print(f"saved adapter with {len(names)} tensors, e.g. {names[:2]}")


def prefix_adapter(src: str, out: str, prefix: str) -> None:
    """Copy an adapter, renaming `base_model.model.X` to `base_model.model.<prefix>X`.

    Mimics an adapter trained on e.g. `BertForMaskedLM` (weights under `bert.`).
    """
    import shutil

    from safetensors.torch import load_file, save_file

    shutil.copytree(src, out, dirs_exist_ok=True)
    tensors = load_file(f"{src}/adapter_model.safetensors")
    head = "base_model.model."
    renamed = {head + prefix + k[len(head) :]: v for k, v in tensors.items()}
    save_file(renamed, f"{out}/adapter_model.safetensors")
    print(f"prefixed {len(renamed)} tensors, e.g. {sorted(renamed)[0]}")


def hf_embeddings(model_name, adapter_dir, pooling, normalize):
    from peft import PeftModel
    from transformers import AutoModel, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float32).eval()
    if adapter_dir:
        model = PeftModel.from_pretrained(model, adapter_dir).merge_and_unload()
    out = []
    with torch.inference_mode():
        for text in TEXTS:
            enc = tok(text, return_tensors="pt")
            hidden = model(**enc).last_hidden_state[0]
            vec = hidden[0] if pooling == "CLS" else hidden.mean(dim=0)
            out.append(F.normalize(vec, dim=-1) if normalize else vec)
    return torch.stack(out)


def check(model_name, adapter_dir, hf_adapter_dir, expect_unsupported) -> int:
    from vllm import LLM
    from vllm.lora.request import LoRARequest

    result = {"model": model_name, "adapter": adapter_dir}
    try:
        llm = LLM(
            model=model_name,
            runner="pooling",
            dtype="float16",  # LoRA shrink kernel needs fp16/bf16; T4 has no bf16
            enforce_eager=True,
            max_model_len=256,
            enable_lora=True,
            max_lora_rank=8,
            gpu_memory_utilization=0.5,
        )
    except Exception as e:  # noqa: BLE001
        msg = f"{type(e).__name__}: {e}"
        result.update(status="lora_rejected", error=msg[:400])
        print("RESULT", json.dumps(result))
        return 0 if expect_unsupported else 1

    if expect_unsupported:
        result["status"] = "unexpectedly_accepted"
        print("RESULT", json.dumps(result))
        return 1

    pooler = llm.llm_engine.vllm_config.model_config.pooler_config
    pooling = pooler.seq_pooling_type
    lora = LoRARequest("random-adapter", 1, adapter_dir)
    v_base = torch.tensor([o.outputs.embedding for o in llm.embed(TEXTS)])
    v_lora = torch.tensor(
        [o.outputs.embedding for o in llm.embed(TEXTS, lora_request=lora)]
    )
    h_base = hf_embeddings(model_name, None, pooling, normalize=True)
    h_lora = hf_embeddings(model_name, hf_adapter_dir, pooling, normalize=True)

    cos = lambda a, b: F.cosine_similarity(a, b, dim=-1)  # noqa: E731
    base_match = cos(v_base, h_base).min().item()
    lora_match = cos(v_lora, h_lora).min().item()
    lora_effect = cos(v_lora, v_base).max().item()
    lora_vs_hf_base = cos(v_lora, h_base).max().item()
    result.update(
        pooling=pooling,
        min_cos_vllm_base_vs_hf_base=round(base_match, 6),
        min_cos_vllm_lora_vs_hf_lora=round(lora_match, 6),
        max_cos_vllm_lora_vs_vllm_base=round(lora_effect, 6),
        max_cos_vllm_lora_vs_hf_base=round(lora_vs_hf_base, 6),
    )
    ok = base_match > 0.995 and lora_match > 0.995 and lora_effect < 0.99
    result["status"] = "pass" if ok else "fail"
    print("RESULT", json.dumps(result))
    return 0 if ok else 1


if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "make-adapter":
        make_adapter(sys.argv[2], sys.argv[3])
    elif cmd == "prefix-adapter":
        prefix_adapter(sys.argv[2], sys.argv[3], sys.argv[4])
    elif cmd == "check":
        args = sys.argv
        hf_dir = args[args.index("--hf-adapter") + 1] if "--hf-adapter" in args else args[3]
        sys.exit(check(args[2], args[3], hf_dir, "--expect-unsupported" in args))
    else:
        raise SystemExit(__doc__)