import argparse
import secrets

from huggingface_hub import HfApi, Volume


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Provision HF Space persistence/settings for the RCA Label Studio Space."
    )
    parser.add_argument("--space-id", required=True, help="Space repo id, e.g. namespace/space-name.")
    parser.add_argument("--bucket-id", required=True, help="HF bucket id, e.g. namespace/label-studio-data.")
    parser.add_argument("--secret-key", default="", help="Stable Django SECRET_KEY. Generated if omitted.")
    parser.add_argument("--admin-email", default="", help="Optional Label Studio admin email.")
    parser.add_argument("--admin-password", default="", help="Optional Label Studio admin password.")
    parser.add_argument("--admin-token", default="", help="Optional stable Label Studio legacy API token.")
    parser.add_argument("--factory-reboot", action="store_true", help="Factory reboot after provisioning.")
    args = parser.parse_args()

    api = HfApi()
    secret_key = args.secret_key or secrets.token_urlsafe(48)
    admin_email = args.admin_email or ""
    admin_password = args.admin_password or ""
    admin_token = args.admin_token or ""
    api.create_bucket(args.bucket_id, private=True, exist_ok=True)

    api.set_space_volumes(
        args.space_id,
        volumes=[
            Volume(type="bucket", source=args.bucket_id, mount_path="/data"),
        ],
    )

    api.add_space_variable(args.space_id, "LABEL_STUDIO_BASE_DATA_DIR", "/data")
    api.add_space_variable(args.space_id, "STORAGE_PERSISTENCE", "1")
    api.add_space_variable(args.space_id, "LABEL_STUDIO_ENABLE_LEGACY_API_TOKEN", "true")
    api.add_space_variable(args.space_id, "LABEL_STUDIO_DISABLE_SIGNUP_WITHOUT_LINK", "true")

    api.add_space_secret(args.space_id, "SECRET_KEY", secret_key)
    if admin_email:
        api.add_space_secret(args.space_id, "LABEL_STUDIO_USERNAME", admin_email)
    if admin_password:
        api.add_space_secret(args.space_id, "LABEL_STUDIO_PASSWORD", admin_password)
    if admin_token:
        api.add_space_secret(args.space_id, "LABEL_STUDIO_USER_TOKEN", admin_token)

    if args.factory_reboot:
        api.restart_space(args.space_id, factory_reboot=True)

    print(f"Provisioned {args.space_id} with bucket {args.bucket_id} mounted at /data.")


if __name__ == "__main__":
    main()