defmodule Mix.Tasks.Hf.Download do
  @shortdoc "Download a file or entire repo from the HuggingFace Hub"
  @moduledoc """
  Downloads a file or an entire repository snapshot from the HuggingFace Hub.

      $ mix hf.download gpt2
      $ mix hf.download gpt2 config.json
      $ mix hf.download gpt2 config.json --revision main --local-dir ./my-model
      $ mix hf.download --type dataset rajpurkar/squad

  ## Options

    * `--type` / `-t` — repo type: `model`, `dataset`, `space` (default: `model`)
    * `--revision` / `-r` — branch or commit (default: `main`)
    * `--local-dir` — save to this local directory instead of cache
    * `--token` — HF API token (uses saved token if omitted)
    * `--quiet` — suppress progress output
  """

  use Mix.Task

  @impl Mix.Task
  def run(args) do
    {opts, argv, _} =
      OptionParser.parse(args,
        aliases: [t: :type, r: :revision, q: :quiet],
        strict: [type: :string, revision: :string, local_dir: :string, token: :string, quiet: :boolean]
      )

    [repo_id | rest] =
      case argv do
        [] -> Mix.raise("Usage: mix hf.download REPO_ID [FILENAME]")
        parts -> parts
      end

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

    case List.first(rest) do
      nil -> download_snapshot(repo_id, opts, token)
      filename -> download_single(repo_id, filename, opts, token)
    end
  end

  defp download_single(repo_id, filename, opts, token) do
    unless opts[:quiet], do: Mix.shell().info("Downloading #{repo_id}/#{filename}...")

    dl_opts = [
      filename: filename,
      revision: opts[:revision] || "main",
      local_dir: opts[:local_dir],
      repo_type: String.to_existing_atom(opts[:type] || "model"),
      access_token: token
    ]

    case HuggingfaceClient.Hub.Files.hf_hub_download(repo_id, dl_opts) do
      {:ok, path} -> Mix.shell().info("✓ Saved to: #{path}")
      {:error, e} -> Mix.raise("Download failed: #{Exception.message(e)}")
    end
  end

  defp download_snapshot(repo_id, opts, token) do
    unless opts[:quiet], do: Mix.shell().info("Downloading snapshot of #{repo_id}...")

    dl_opts = [
      revision: opts[:revision] || "main",
      local_dir: opts[:local_dir],
      repo_type: String.to_existing_atom(opts[:type] || "model"),
      access_token: token
    ]

    case HuggingfaceClient.Hub.Snapshots.snapshot_download(repo_id, dl_opts) do
      {:ok, dir} -> Mix.shell().info("✓ Snapshot saved to: #{dir}")
      {:error, e} -> Mix.raise("Snapshot failed: #{Exception.message(e)}")
    end
  end
end