"""
Run this ON THE MACHINE WHERE THE CHECKPOINTS ACTUALLY LIVE (the cluster) --
these paths (/scratch/uzair/..., /datasets/uzair/...) are not reachable from
a laptop or from the HF Space itself.

Uploads to two separate Hub repos, matching MODEL_REPO/TOKENIZERS_REPO in
hf_space_demo/inference.py:
  - the main 4M checkpoint  -> --model_repo_id      (e.g. EPFL-VILAB/Video-4M-models)
  - the 7 tokenizer ckpts   -> --tokenizers_repo_id (e.g. EPFL-VILAB/Video-4M-tokenizers),
    one folder per modality (rgb/, surface-normals/, depth/, opticalflow/, v-jepa-2/, dinov2/, siglip-2/)

They're split because the model checkpoint changes often during training
while the tokenizers are stable, so each gets its own upload/version history.

Usage:
    python upload_checkpoints_to_hub.py \\
        --model_repo_id EPFL-VILAB/Video-4M-models \\
        --tokenizers_repo_id EPFL-VILAB/Video-4M-tokenizers \\
        --model_path /scratch/uzair/A18_large_130_checkpoint.pth \\
        --vidtok_rgb /scratch/uzair/clariden/vidok_rgb_tokenizer/rgb_normal_last_ckpt_580000_iter.ckpt \\
        --vidtok_normal /scratch/uzair/clariden/vidtok_normal_tokenizer/second_run_lr_by_2/420000.ckpt \\
        --vidtok_depth /scratch/uzair/clariden/vidtok_depth_tokenizers/depth_first_run_lsat_ckpt_500000.ckpt \\
        --vidtok_opticalflow /scratch/uzair/clariden/vidtok_optical_flow_tokenizer_weights/first_run_full_finetuning.ckpt \\
        --vjepa /datasets/uzair/weights_from_clariden/vjepa_fixed_tokenizer_weights/.../epoch=3-step=150000.ckpt \\
        --dinov2 /scratch/uzair/clariden/tokenizers_for_feature_maps_michael/tokenizers/checkpoints/dinov2_l1_176_411_16807.ckpt \\
        --siglipv2 /scratch/uzair/clariden/tokenizers_for_feature_maps_michael/tokenizers/checkpoints/siglip_ens_176_411_16807.ckpt

Pass --model_only or --tokenizers_only to upload just one side (e.g. once
the tokenizers are up, you'll only need --model_only on later runs whenever
the model checkpoint changes).
"""
import argparse

from huggingface_hub import HfApi, create_repo

# Must match MODEL_WEIGHT_FILES / TOKENIZER_WEIGHT_FILES in hf_space_demo/inference.py
MODEL_DEST_FILENAMES = {
    "model_path": "main_model/checkpoint.pth",
}
TOKENIZER_DEST_FILENAMES = {
    "vidtok_rgb": "rgb/ckpt.ckpt",
    "vidtok_normal": "surface-normals/ckpt.ckpt",
    "vidtok_depth": "depth/ckpt.ckpt",
    "vidtok_opticalflow": "opticalflow/ckpt.ckpt",
    "vjepa": "v-jepa-2/ckpt.ckpt",
    "dinov2": "dinov2/ckpt.ckpt",
    "siglipv2": "siglip-2/ckpt.ckpt",
}


def upload(api, repo_id, private, local_paths_by_key, dest_filenames):
    create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
    for key, dest_filename in dest_filenames.items():
        local_path = local_paths_by_key[key]
        print(f"Uploading {local_path} -> {repo_id}:{dest_filename}")
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=dest_filename,
            repo_id=repo_id,
            repo_type="model",
        )


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model_repo_id", default="EPFL-VILAB/Video-4M-models")
    parser.add_argument("--tokenizers_repo_id", default="EPFL-VILAB/Video-4M-tokenizers")
    parser.add_argument("--private", action="store_true", default=True, help="Create repos as private (default)")
    parser.add_argument("--public", dest="private", action="store_false")
    parser.add_argument("--model_only", action="store_true")
    parser.add_argument("--tokenizers_only", action="store_true")

    parser.add_argument("--model_path", help="Local path to the main 4M checkpoint")
    for key in TOKENIZER_DEST_FILENAMES:
        parser.add_argument(f"--{key}", help="Local path to this tokenizer checkpoint on this machine")
    args = parser.parse_args()

    do_model = not args.tokenizers_only
    do_tokenizers = not args.model_only
    if do_model and not args.model_path:
        parser.error("--model_path is required unless --tokenizers_only is set")
    if do_tokenizers and any(getattr(args, key) is None for key in TOKENIZER_DEST_FILENAMES):
        missing = [key for key in TOKENIZER_DEST_FILENAMES if getattr(args, key) is None]
        parser.error(f"missing tokenizer paths (or pass --model_only): {missing}")

    api = HfApi()

    if do_model:
        upload(api, args.model_repo_id, args.private, {"model_path": args.model_path}, MODEL_DEST_FILENAMES)
    if do_tokenizers:
        tokenizer_paths = {key: getattr(args, key) for key in TOKENIZER_DEST_FILENAMES}
        upload(api, args.tokenizers_repo_id, args.private, tokenizer_paths, TOKENIZER_DEST_FILENAMES)

    print("\nDone. Set these env vars where the Space/app runs:")
    if do_model:
        print(f"  FOURM_MODEL_REPO={args.model_repo_id}")
    if do_tokenizers:
        print(f"  FOURM_TOKENIZERS_REPO={args.tokenizers_repo_id}")


if __name__ == "__main__":
    main()