"""Obtain the six external inference assets, without downloading datasets."""
from pathlib import Path
import os

UPSTREAM_COMMIT = "4faebb83ccda8630e27fbe2189aed1928f4623a0"
SALAD_DRIVE_FOLDER = "1YuDQCgc6RJ4WlR9vt_L34nkXChC7C98_"
ASSETS = (
    "t2m/t2m_denoiser_vpred_vaegelu/opt.txt",
    "t2m/t2m_denoiser_vpred_vaegelu/model/net_best_fid.tar",
    "t2m/t2m_vae_gelu/opt.txt",
    "t2m/t2m_vae_gelu/model/net_best_fid.tar",
    "t2m/Comp_v6_KLD005/meta/mean.npy",
    "t2m/Comp_v6_KLD005/meta/std.npy",
)


def ensure_assets(source_root, report=lambda message: None):
    source_root = Path(source_root)
    checkpoint = source_root / "checkpoints/t2sm/ours/epoch_0100.ckpt"
    if not checkpoint.is_file() or checkpoint.stat().st_size == 0:
        raise RuntimeError("Style-SALAD checkpoint missing. Use the complete upstream checkout.")
    root = source_root / "checkpoints"
    missing = [name for name in ASSETS if not (root / name).is_file()
               or (root / name).stat().st_size == 0]
    if not missing:
        return
    if os.environ.get("AUTO_DOWNLOAD_ASSETS", "1") != "1":
        raise RuntimeError("Missing assets: " + ", ".join(missing))

    # Optional mirror owned by the user, with paths t2m/... at repo root.
    # Credentials are read by the Hub client, never from browser fields.
    mirror = os.environ.get("ASSET_REPO_ID", "").strip()
    if mirror:
        from huggingface_hub import hf_hub_download
        import shutil
        for name in missing:
            report("Obtaining model weights and statistics...")
            cached = hf_hub_download(repo_id=mirror, filename=name,
                                    token=os.environ.get("HF_TOKEN") or None)
            target = root / name
            target.parent.mkdir(parents=True, exist_ok=True)
            temporary = target.with_name(target.name + ".download")
            shutil.copyfile(cached, temporary)
            temporary.replace(target)
    else:
        import gdown
        report("Reading the published SALAD asset list...")
        entries = gdown.download_folder(id=SALAD_DRIVE_FOLDER,
                                        skip_download=True, quiet=True)
        if not entries:
            raise RuntimeError("Could not list the published weights on Google Drive.")
        for name in missing:
            suffix = name.removeprefix("t2m/")
            candidates = [e for e in entries
                          if e.path.replace("\\", "/") == suffix
                          or e.path.replace("\\", "/").endswith("/" + suffix)]
            if len(candidates) != 1:
                raise RuntimeError("Asset not uniquely identified on Drive: " + name)
            target = root / name
            target.parent.mkdir(parents=True, exist_ok=True)
            temporary = target.with_name(target.name + ".download")
            report("Downloading " + Path(name).name + "...")
            result = gdown.download(id=candidates[0].id, output=str(temporary), quiet=False)
            if not result or not temporary.is_file() or not temporary.stat().st_size:
                raise RuntimeError("Incomplete download: " + name)
            temporary.replace(target)
    missing = [name for name in ASSETS if not (root / name).is_file()
               or not (root / name).stat().st_size]
    if missing:
        raise RuntimeError("Assets still missing: " + ", ".join(missing))