"""Pin/download checkpoints and adapt Transformers 5 tokenizer metadata to v4."""
import os
import json
from pathlib import Path
from huggingface_hub import snapshot_download

SPECS = [
    ("intfloat/multilingual-e5-small",
     "fd1525a9fd15316a2d503bf26ab031a61d056e98",
     "E5_MODEL_DIR", "/opt/models/multilingual-e5-small"),
    ("Qwen/Qwen3-Embedding-0.6B",
     "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
     "QWEN_MODEL_DIR", "/opt/models/qwen3-embedding-0.6b"),
    ("lightonai/mLateOn",
     "edd378f",
     "COLBERT_MODEL_DIR", "/opt/models/mlateon"),
    ("BAAI/bge-m3",
     "5617a9f",
     "BGE_M3_MODEL_DIR", "/opt/models/bge-m3"),
]
def adapt_mlateon_tokenizer(model_dir: str | Path) -> None:
    """Transformers v5 calls the generic fast tokenizer TokenizersBackend;
    Transformers v4 calls it PreTrainedTokenizerFast. All vocabulary and
    tokenizer.json bytes stay unchanged.
    """
    folder = Path(model_dir)
    config_path = folder / "tokenizer_config.json"
    config = json.loads(config_path.read_text(encoding="utf-8"))
    cls = config.get("tokenizer_class")
    if cls != "TokenizersBackend":
        raise RuntimeError(f"Unexpected mLateOn tokenizer_class {cls!r}")
    if not (folder / "tokenizer.json").is_file():
        raise RuntimeError("mLateOn tokenizer.json not found")
    config["tokenizer_class"] = "PreTrainedTokenizerFast"
    config_path.write_text(
        json.dumps(config, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )


for repo_id, revision, env_name, default in SPECS:
    # BGE-M3 has NO safetensors encoder: pytorch_model.bin and the
    # sparse_linear.pt / colbert_linear.pt heads are all mandatory.
    ignored = ["*.onnx", "*.h5", "*.msgpack", "*.ot",
               "*.md", "*.tflite", "*.gguf"]
    if repo_id != "BAAI/bge-m3":
        ignored += ["*.bin", "*.pt"]
    path = snapshot_download(
        repo_id=repo_id, revision=revision,
        local_dir=os.getenv(env_name, default),
        ignore_patterns=ignored,
    )
    if repo_id == "BAAI/bge-m3":
        for filename in ("pytorch_model.bin", "sparse_linear.pt",
                         "colbert_linear.pt"):
            if not (Path(path) / filename).is_file():
                raise RuntimeError(f"Missing BGE-M3 tri-head weight: {filename}")
    print(f"Downloaded {repo_id}@{revision} to {path}", flush=True)

    if repo_id == "lightonai/mLateOn":
        adapt_mlateon_tokenizer(path)
        # Fail the Docker build before release if the real tokenizer fails.
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True)
        if not tokenizer("Hello שלום привет", truncation=False)["input_ids"]:
            raise RuntimeError("mLateOn tokenizer produced no tokens")
        print("mLateOn fast-tokenizer build preflight passed", flush=True)