import os
from pathlib import Path
from huggingface_hub import snapshot_download

def get_weights_dir(repo_id: str) -> Path:
    # repo_id must be like "BiasLab2025/taskclip-weights" (NOT a URL)
    repo_id = repo_id.strip()
    if repo_id.startswith("http"):
        # allow passing a full URL by accident
        repo_id = repo_id.rstrip("/").split("huggingface.co/")[-1]

    token = os.getenv("HF_TOKEN")  # only needed if the repo is private

    p = snapshot_download(
        repo_id=repo_id,
        repo_type="model",          # IMPORTANT for your weights repo
        local_dir="weights_cache",
        local_dir_use_symlinks=False,
        token=token,
    )
    return Path(p).resolve()