defmodule Mix.Tasks.Hf.Providers do
  @shortdoc "List all registered HuggingFace Inference providers and their supported tasks"

  @moduledoc """
  Lists all providers registered in `HuggingfaceClient.Inference.ProviderRegistry`.

      $ mix hf.providers

      Provider               Tasks
      ─────────────────────────────────────────────────────
      baseten                conversational
      black-forest-labs      text-to-image
      cerebras               conversational
      clarifai               conversational
      cohere                 conversational
      deepinfra              conversational, text-generation
      fal-ai                 automatic-speech-recognition, image-segmentation, ...
      ...

  ## Options

      --provider PROVIDER   Show only tasks for a specific provider
      --json                Output as JSON
  """

  use Mix.Task

  @impl Mix.Task
  def run(args) do
    {opts, _, _} =
      OptionParser.parse(args, strict: [provider: :string, json: :boolean])

    pairs = HuggingfaceClient.Inference.ProviderRegistry.all()
    filtered = if opts[:provider], do: Enum.filter(pairs, fn {p, _} -> p == opts[:provider] end), else: pairs

    grouped =
      filtered
      |> Enum.group_by(fn {provider, _} -> provider end, fn {_, task} -> task end)
      |> Enum.map(fn {k, v} -> {k, Enum.sort(v)} end)
      |> Enum.sort_by(fn {k, _} -> k end)

    if opts[:json] do
      grouped
      |> Enum.into(%{})
      |> Jason.encode!(pretty: true)
      |> Mix.shell().info()
    else
      Mix.shell().info("\nProvider               Tasks")
      Mix.shell().info(String.duplicate("─", 72))

      Enum.each(grouped, fn {provider, tasks} ->
        pad = String.pad_trailing(provider, 22)
        joined = Enum.join(tasks, ", ")
        Mix.shell().info("#{pad} #{joined}")
      end)

      Mix.shell().info("\nTotal: #{length(grouped)} providers, #{length(pairs)} task registrations\n")
    end
  end
end