defmodule Mix.Tasks.Hf.ModelInfo do
  @shortdoc "Fetch model info and available providers from the HuggingFace Hub"

  @moduledoc """
  Fetches metadata for a model from the HuggingFace Hub and displays it,
  including all available inference providers.

      $ mix hf.model_info meta-llama/Llama-3.1-8B-Instruct

      Model:         meta-llama/Llama-3.1-8B-Instruct
      Task:          text-generation
      Library:       transformers
      Downloads:     1_234_567
      Likes:         8_901
      Gated:         false
      Private:       false

      Available providers (status: live)
      ─────────────────────────────────
      groq             conversational   llama-3.1-8b-instant
      together         conversational   meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo
      nebius           conversational   meta-llama/Meta-Llama-3.1-8B-Instruct
      ...

  ## Options

      --token TOKEN    HuggingFace access token (or set HF_TOKEN env var)
      --json           Output as JSON
  """

  use Mix.Task

  @impl Mix.Task
  def run(args) do
    Application.ensure_all_started(:huggingface_client)
    Application.ensure_all_started(:req)

    {opts, [model_id | _], _} =
      OptionParser.parse(args, strict: [token: :string, json: :boolean])

    token = opts[:token] || System.get_env("HF_TOKEN")

    Mix.shell().info("Fetching info for #{model_id}...")

    case HuggingfaceClient.Inference.ModelInfo.fetch(model_id, access_token: token) do
      {:ok, info} ->
        if opts[:json] do
          info
          |> Map.from_struct()
          |> Jason.encode!(pretty: true)
          |> Mix.shell().info()
        else
          print_info(info)
        end

      {:error, err} ->
        Mix.shell().error("Error: #{err.message}")
        exit({:shutdown, 1})
    end
  end

  defp print_info(info) do
    Mix.shell().info("")
    Mix.shell().info("Model:         #{info.id}")
    Mix.shell().info("Task:          #{info.pipeline_tag || "unknown"}")
    Mix.shell().info("Library:       #{info.library_name || "unknown"}")
    Mix.shell().info("Downloads:     #{format_number(info.downloads)}")
    Mix.shell().info("Likes:         #{format_number(info.likes)}")
    Mix.shell().info("Gated:         #{info.gated}")
    Mix.shell().info("Private:       #{info.private}")

    if info.tags != [] do
      Mix.shell().info("Tags:          #{Enum.join(Enum.take(info.tags, 10), ", ")}")
    end

    live = Enum.filter(info.providers, &(&1["status"] == "live"))

    if live == [] do
      Mix.shell().info("\nNo inference providers available for this model.\n")
    else
      Mix.shell().info("\nAvailable providers (status: live)")
      Mix.shell().info(String.duplicate("─", 56))

      Enum.each(live, fn p ->
        provider = String.pad_trailing(p["provider"] || "", 16)
        task = String.pad_trailing(p["task"] || "", 16)
        pid = p["provider_id"] || ""
        Mix.shell().info("#{provider} #{task} #{pid}")
      end)

      Mix.shell().info("")
    end
  end

  defp format_number(nil), do: "unknown"
  defp format_number(n) when n >= 1_000_000, do: "#{Float.round(n / 1_000_000, 1)}M"
  defp format_number(n) when n >= 1_000, do: "#{Float.round(n / 1_000, 1)}K"
  defp format_number(n), do: "#{n}"
end