exmobiledevice

defmodule ExMobileDevice.Diagnostics.RestartManager do
  @moduledoc false
  use GenServer, restart: :temporary
  use TypedStruct

  alias ExMobileDevice.Diagnostics.RestartSupervisor
  alias ExMobileDevice.{Services, TaskSupervisor}

  require Logger

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

  @doc """
  Makes and keeps an SSL connection to the Diagnostics service of a device for later use
  """
  @spec preconnect(binary()) :: :ok | {:error, any()}
  def preconnect(udid) do
    child_spec = {__MODULE__, udid: udid}

    with {:ok, _pid} <- DynamicSupervisor.start_child(RestartSupervisor, child_spec) do
      :ok
    end
  end

  @doc """
  Sends a request to the Diagnostics service using the preconnected SSL socket
  """
  @spec request(binary(), binary(), timeout()) :: :ok | {:error, any()}
  def request(udid, request, timeout \\ @default_tmo) do
    GenServer.call(via(udid), {:request, request}, timeout)
  catch
    # Preconnect might have failed
    :exit, {:noproc, _} ->
      {:error, :failed}
  end

  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(args) do
    udid = Keyword.fetch!(args, :udid)
    GenServer.start_link(__MODULE__, args, name: via(udid))
  end

  typedstruct do
    field :udid, String.t()
    field :ssl_sock, :ssl.sslsocket() | nil
  end

  @impl GenServer
  def init(args) do
    udid = Keyword.fetch!(args, :udid)

    deadline = System.monotonic_time(:millisecond) + @default_tmo
    {:ok, %__MODULE__{udid: udid}, {:continue, {:preconnect, deadline}}}
  end

  @impl GenServer
  def handle_continue({:preconnect, deadline}, %__MODULE__{ssl_sock: nil} = state) do
    remaining = deadline - System.monotonic_time(:millisecond)

    cp = self()

    case run_in_task(
           fn ->
             with {:ok, ssl_sock} <- Services.connect(state.udid, @diagnostics, remaining) do
               :ok = :ssl.controlling_process(ssl_sock, cp)
               {:ok, ssl_sock}
             end
           end,
           remaining
         ) do
      {:ok, ssl_sock} ->
        Logger.info("#{__MODULE__}: Preconnected #{state.udid}")
        # Set to active mode in order to receive ssl_close
        :ssl.setopts(ssl_sock, active: true)
        {:noreply, %{state | ssl_sock: ssl_sock}}

      {:error, :timeout} ->
        Logger.error("#{__MODULE__}: #{state.udid} Timed out preconnecting socket")
        {:stop, {:shutdown, :timeout}, state}

      error ->
        Logger.error(
          "#{__MODULE__}: Error preconnecting #{state.udid} SSL socket. #{inspect(error)}"
        )

        # Retry until we exceed the deadline
        Process.send_after(self(), {:preconnect, deadline}, @preconnect_retry_wait_ms)
        {:noreply, state}
    end
  end

  def handle_continue({:preconnect, :deadline}, %__MODULE__{} = state) do
    Logger.error("#{__MODULE__}: Preconnect failed for #{state.udid}. Already connected")
    {:noreply, state}
  end

  @impl GenServer
  def handle_call({:request, _}, _from, %__MODULE__{ssl_sock: nil} = state) do
    {:reply, :not_connected, state}
  end

  def handle_call({:request, request}, _from, %__MODULE__{ssl_sock: ssl_sock} = state) do
    Logger.info("#{__MODULE__}: Trying #{request} request for #{state.udid}...")
    :ssl.setopts(ssl_sock, active: false)

    case Services.rpc(ssl_sock, request, %{}) do
      {:ok, %{"Status" => "Success"}} ->
        Logger.info("#{__MODULE__}: #{request} request for #{state.udid} accepted")
        {:stop, {:shutdown, :closed}, :ok, state}

      error ->
        :ssl.setopts(ssl_sock, active: true)
        {:reply, error, state}
    end
  end

  @impl GenServer
  def handle_info({:preconnect, deadline}, %__MODULE__{} = state) do
    {:noreply, state, {:continue, {:preconnect, deadline}}}
  end

  def handle_info({:ssl_closed, ssl_sock}, %__MODULE__{ssl_sock: ssl_sock} = state) do
    Logger.debug("#{__MODULE__}: SSL socket closed. Stopping #{state.udid}")
    {:stop, {:shutdown, :closed}, state}
  end

  def handle_info(_, %__MODULE__{} = state) do
    {:noreply, state}
  end

  defp via(udid), do: {:via, Registry, {ExMobileDevice.Diagnostics.Registry, udid}}

  defp run_in_task(_fun, timeout) when timeout <= 0, do: {:error, :timeout}

  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