defmodule ExMobileDevice.Diagnostics do
  @moduledoc """
  Services for restarting, shutting down, or sleeping a device.
  """

  alias ExMobileDevice.Diagnostics.RestartManager
  alias ExMobileDevice.Services
  alias ExMobileDevice.TaskSupervisor

  require Logger

  @diagnostics "com.apple.mobile.diagnostics_relay"
  @default_tmo 30_000

  @doc """
  Restart the specified device.
  """
  @spec restart(String.t(), timeout()) :: :ok | {:error, any}
  def restart(udid, timeout \\ @default_tmo) do
    run_in_task(udid, "Restart", timeout)
  end

  @doc """
  Shutdown the specified device.
  """
  @spec shutdown(String.t(), timeout()) :: :ok | {:error, any}
  def shutdown(udid, timeout \\ @default_tmo) do
    run_in_task(udid, "Shutdown", timeout)
  end

  @doc """
  Send the specified device to sleep.
  """
  @spec sleep(String.t(), timeout()) :: :ok | {:error, any}
  def sleep(udid, timeout \\ @default_tmo) do
    run_in_task(udid, "Sleep", timeout)
  end

  @doc false
  def ioreg(udid, opts \\ []) do
    args =
      opts
      |> Keyword.take([:current_plane, :entry_name, :entry_class])
      |> Map.new(fn {k, v} -> {to_ioreg_arg(k), v} end)

    timeout = Keyword.get(opts, :timeout, @default_tmo)

    run_in_task(
      fn ->
        with {:ok, ssl_sock} <- Services.connect(udid, @diagnostics) do
          case Services.rpc(ssl_sock, "IORegistry", args) do
            {:ok, %{"Status" => "Success"} = response} ->
              {:ok, get_in(response, ["Diagnostics", "IORegistry"])}

            _ ->
              {:error, :failed}
          end
        end
      end,
      timeout
    )
  end

  defp to_ioreg_arg(:current_plane), do: "CurrentPlane"
  defp to_ioreg_arg(:entry_class), do: "EntryClass"
  defp to_ioreg_arg(:entry_name), do: "EntryName"

  defp run_in_task(udid, request, timeout) when is_binary(request) do
    run_in_task(
      fn ->
        with {:ok, ssl_sock} <- Services.connect(udid, @diagnostics, timeout),
             {:ok, %{"Status" => "Success"}} <- Services.rpc(ssl_sock, request, %{}) do
          :ok
        else
          error ->
            Logger.error(
              "#{__MODULE__}: #{request} request for #{udid} failed. #{inspect(error)}"
            )

            # Try with RestartManager
            Logger.warning("Attempting #{request} request for #{udid} using RestartManager...")

            RestartManager.request(udid, request, timeout)
        end
      end,
      # Should never get here. Give the Services.connect call enough time to timeout
      ceil(timeout * 1.1)
    )
  end

  defp run_in_task(fun, timeout) when is_function(fun, 0) do
    task = TaskSupervisor.async_nolink(fun)

    case Task.yield(task, timeout) || Task.shutdown(task) do
      {:ok, result} ->
        result

      nil ->
        {:error, :timeout}
    end
  end
end