lib/mix/tasks/hf.spaces.ex

defmodule Mix.Tasks.Hf.Spaces do
  @shortdoc "List and manage your HuggingFace Spaces"
  @moduledoc """
  List and manage HuggingFace Spaces.

      $ mix hf.spaces
      $ mix hf.spaces --author my-org --sdk gradio
      $ mix hf.spaces --pause my-org/my-space
      $ mix hf.spaces --restart my-org/my-space

  ## Options

    * `--author` / `-a` — filter by author/org
    * `--sdk` — filter by SDK: `gradio`, `streamlit`, `docker`, `static`
    * `--pause SPACE_ID` — pause a Space
    * `--restart SPACE_ID` — restart a Space
    * `--runtime SPACE_ID` — show Space runtime info
    * `--token` — HF API token
  """

  use Mix.Task

  alias HuggingfaceClient.Config
  alias HuggingfaceClient.Hub.Search
  alias HuggingfaceClient.Hub.Spaces

  @impl Mix.Task
  def run(args) do
    {opts, _, _} =
      OptionParser.parse(args,
        aliases: [a: :author],
        strict: [
          author: :string,
          sdk: :string,
          pause: :string,
          restart: :string,
          runtime: :string,
          token: :string
        ]
      )

    token = opts[:token] || Config.token()

    cond do
      space = opts[:pause] -> do_pause(space, token)
      space = opts[:restart] -> do_restart(space, token)
      space = opts[:runtime] -> do_runtime(space, token)
      true -> do_list(opts, token)
    end
  end

  defp do_pause(space, token) do
    Mix.shell().info("Pausing #{space}...")

    case Spaces.pause(space, access_token: token) do
      {:ok, _} -> Mix.shell().info("✓ Paused")
      err -> Mix.raise("Failed: #{inspect(err)}")
    end
  end

  defp do_restart(space, token) do
    Mix.shell().info("Restarting #{space}...")

    case Spaces.restart(space, access_token: token) do
      {:ok, _} -> Mix.shell().info("✓ Restarted")
      err -> Mix.raise("Failed: #{inspect(err)}")
    end
  end

  defp do_runtime(space, token) do
    case Spaces.get_runtime(space, access_token: token) do
      {:ok, rt} ->
        Mix.shell().info("Space: #{space}")
        Mix.shell().info("Stage:    #{rt["stage"]}")
        Mix.shell().info("Hardware: #{rt["hardware"] || "default"}")
        Mix.shell().info("SDK:      #{rt["sdk"] || "unknown"}")

      err ->
        Mix.raise("Failed: #{inspect(err)}")
    end
  end

  defp do_list(opts, token) do
    search_opts =
      [access_token: token]
      |> maybe_add(opts, :author)
      |> maybe_add(opts, :sdk)

    spaces = Enum.take(Search.spaces(search_opts), 20)

    if spaces == [] do
      Mix.shell().info("No Spaces found.")
    else
      print_spaces_table(spaces)
    end
  end

  defp print_spaces_table(spaces) do
    Mix.shell().info(String.pad_trailing("ID", 50) <> "  SDK      Likes")
    Mix.shell().info(String.duplicate("─", 72))

    Enum.each(spaces, fn s ->
      id = s["id"] || ""
      sdk = String.pad_trailing(s["sdk"] || "-", 8)
      likes = s["likes"] || 0
      Mix.shell().info("#{String.pad_trailing(id, 50)}  #{sdk}  #{likes}")
    end)
  end

  defp maybe_add(opts_kw, opts, key) do
    if opts[key], do: Keyword.put(opts_kw, key, opts[key]), else: opts_kw
  end
end