#!/usr/bin/env python3
import argparse
import os
from typing import Optional

from huggingface_hub import HfApi, SpaceStorage


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Request (or delete) persistent storage for a Hugging Face Space.")
    parser.add_argument("repo_id", help="Space repository ID in the form 'username/space-name'")
    parser.add_argument(
        "--storage",
        "-s",
        choices=[e.name for e in SpaceStorage],
        default=SpaceStorage.SMALL.name,
        help="Storage tier to request (default: SMALL)",
    )
    parser.add_argument(
        "--delete",
        action="store_true",
        help="Delete the Space storage instead of requesting it",
    )
    parser.add_argument(
        "--token",
        "-t",
        default=None,
        help="Hugging Face access token (uses HF_TOKEN or HUGGINGFACE_HUB_TOKEN if not provided)",
    )
    return parser.parse_args()


def resolve_token(token_arg: Optional[str]) -> Optional[str]:
    return token_arg or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")


def main() -> None:
    args = parse_args()
    token = resolve_token(args.token)

    if not token:
        print(
            "Warning: No token provided via --token or environment variable. "
            "Please set HF_TOKEN/HUGGINGFACE_HUB_TOKEN or run `huggingface-cli login`."
        )

    api = HfApi(token=token)

    if args.delete:
        result = api.delete_space_storage(repo_id=args.repo_id)
        print(result)
        return

    storage = SpaceStorage(args.storage)
    result = api.request_space_storage(repo_id=args.repo_id, storage=storage)
    print(result)


if __name__ == "__main__":
    main()