#!/usr/bin/env python3
"""Pause or downgrade the HF training Space after artifacts are retrieved."""

from __future__ import annotations

import argparse

from huggingface_hub import HfApi


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Pause or downgrade the PolyGuard training Space.")
    parser.add_argument("--repo-id", default="TheJackBright/polyguard-openenv-training")
    parser.add_argument("--mode", choices=["pause", "cpu-basic"], default="cpu-basic")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    api = HfApi()
    if args.mode == "pause":
        runtime = api.pause_space(repo_id=args.repo_id)
    else:
        runtime = api.request_space_hardware(repo_id=args.repo_id, hardware="cpu-basic")
    print(f"space={args.repo_id} mode={args.mode} runtime={runtime}")


if __name__ == "__main__":
    main()