exmobiledevice

defmodule ExMobileDevice.WebInspector do
  @moduledoc """
  Functions for controlling Mobile Safari.

  This module provides a means of controlling Safari, without requiring SafariDriver.
  Starting an instance of this process will create a connection to an automated safari
  instance, on the specified device.

  ## Example
      iex(1)> alias ExMobileDevice.WebInspector
      ExMobileDevice.WebInspector
      iex(2)> {:ok, pid} = WebInspector.start_supervised("00008120-0018DEADC0DEFACE")
      {:ok, #PID<0.208.0>}
      iex(3)> {:ok, page} = WebInspector.create_page(pid)
      {:ok, "page-7102B011-5BC0-4785-87DF-ADBA671EAD74"}
      iex(4)> WebInspector.navigate_to(pid, page, "https://elixir-lang.org")
      :ok
  """
  @behaviour :gen_statem
  use TypedStruct

  alias Uniq.UUID
  require Logger

  @doc false
  @spec child_spec(any()) :: Supervisor.child_spec()
  def child_spec(args) do
    %{
      id: __MODULE__,
      start: {__MODULE__, :start_link, [args]},
      restart: :temporary
    }
  end

  @doc """
  Create a supervised connection to the 'webinspector' service on the device.

  The connection is supervised by an internal supervision tree of `exmobiledevice`.
  To manually stop the connection, call `stop/1`.

  > #### Controlling Process {: .info}
  >
  > On success, the caller is the 'controlling process' of the connection. Should
  > the caller exit for any reason, the connection will also exit, which will
  > in turn tear down the automation pages.
  >
  > To transfer control to another process, call `set_controlling_process/2`.
  """
  @spec start_supervised(String.t()) :: DynamicSupervisor.on_start_child()
  def start_supervised(udid) when is_binary(udid) do
    args = [udid: udid, controlling_process: self()]
    DynamicSupervisor.start_child(ExMobileDevice.WebInspector.Supervisor, {__MODULE__, args})
  end

  @doc """
  Stop the webinspector process.

  Stopping the process will terminate the connection to the webinspector service
  on the device, which will remove any automation pages from Safari.
  """
  @spec stop(pid) :: :ok
  def stop(pid) do
    :gen_statem.stop(pid)
  end

  @doc """
  Transfer ownership to the specified process.

  Transfers controlling ownership of `pid` to `cp`. The caller must be the
  current owner.
  """
  @spec set_controlling_process(:gen_statem.server_ref(), pid()) :: :ok | {:error, any()}
  def set_controlling_process(pid, cp) when is_pid(cp) do
    :gen_statem.call(pid, {:set_cp, cp})
  end

  @doc """
  Wait for the session to be created.
  """
  @spec wait_for_session(:gen_statem.server_ref(), non_neg_integer() | :infinity) ::
          :ok | {:error, :failed | :timeout | :running}
  def wait_for_session(pid, timeout \\ 5000) do
    deadline =
      if timeout == :infinity do
        timeout
      else
        :erlang.monotonic_time(:millisecond) + timeout
      end

    :gen_statem.call(pid, {:wait_for_session, deadline})
  end

  @doc """
  Create a new automation page.
  """
  @spec create_page(:gen_statem.server_ref()) :: {:ok, String.t()} | {:error, any()}
  def create_page(pid) do
    :gen_statem.call(pid, :create_page)
  end

  @doc """
  List the current automation pages.
  """
  @spec list_pages(:gen_statem.server_ref()) :: {:ok, list(map())} | {:error, any()}
  def list_pages(pid) do
    :gen_statem.call(pid, :list_pages)
  end

  @doc """
  Switch to the specified page.
  """
  @spec switch_to_page(:gen_statem.server_ref(), String.t()) :: :ok | {:error, any()}
  def switch_to_page(pid, page) do
    :gen_statem.call(pid, {:switch_to_page, page})
  end

  @doc """
  Navigate the specified page to the provided url.

  Supported options are:

  - `timeout`: The page-load timeout in milliseconds. Defaults to 30_000 (30 seconds)
  """
  @spec navigate_to(:gen_statem.server_ref(), String.t(), String.t(), Keyword.t()) ::
          :ok | {:error, any()}
  def navigate_to(pid, page, url, opts \\ []) do
    :gen_statem.call(pid, {:navigate_to, page, url, opts})
  end

  @doc """
  Go to the previous url in the page's history.
  """
  @spec go_back(:gen_statem.server_ref(), String.t()) :: :ok | {:error, any()}
  def go_back(pid, page) do
    :gen_statem.call(pid, {:go_back, page})
  end

  @doc """
  Go to the next url in the page's history.
  """
  @spec go_forward(:gen_statem.server_ref(), String.t()) :: :ok | {:error, any()}
  def go_forward(pid, page) do
    :gen_statem.call(pid, {:go_forward, page})
  end

  @doc """
  Reload the page's current url.
  """
  @spec reload(:gen_statem.server_ref(), String.t()) :: :ok | {:error, any()}
  def reload(pid, page) do
    :gen_statem.call(pid, {:reload, page})
  end

  @doc """
  Take a screenshot of the current page, returning the bytes in PNG format.
  """
  @spec take_screenshot(:gen_statem.server_ref(), String.t()) :: {:ok, binary()} | {:error, any()}
  def take_screenshot(pid, page) do
    :gen_statem.call(pid, {:take_screenshot, page})
  end

  @doc """
  Close the specified page.
  """
  @spec close_page(:gen_statem.server_ref(), String.t()) :: :ok | {:error, any()}
  def close_page(pid, page) do
    :gen_statem.call(pid, {:close_page, page})
  end

  @doc false
  @spec start_link(any()) :: :gen_statem.start_ret()
  def start_link(args) do
    :gen_statem.start_link(__MODULE__, args, hibernate_after: 15_000)
  end

  @service "com.apple.webinspector"
  @bundle "com.apple.mobilesafari"
  @selector "__selector"
  @argument "__argument"

  @rpc_report_identifier "_rpc_reportIdentifier:"
  @rpc_report_current_state "_rpc_reportCurrentState:"
  @rpc_report_connected_applications "_rpc_reportConnectedApplicationList:"
  @rpc_application_connected "_rpc_applicationConnected:"
  @rpc_application_updated "_rpc_applicationUpdated:"
  @rpc_application_disconnected "_rpc_applicationDisconnected:"
  @rpc_application_sent_listing "_rpc_applicationSentListing:"
  @rpc_application_sent_data "_rpc_applicationSentData:"
  @rpc_request_application_launch "_rpc_requestApplicationLaunch:"
  @rpc_forward_automation_session_request "_rpc_forwardAutomationSessionRequest:"
  @rpc_forward_socket_setup "_rpc_forwardSocketSetup:"
  @rpc_forward_socket_data "_rpc_forwardSocketData:"

  @wir_connection_identifier_key "WIRConnectionIdentifierKey"
  @wir_automation_availability_key "WIRAutomationAvailabilityKey"
  @wir_automation_available "WIRAutomationAvailabilityAvailable"
  @wir_application_dictionary_key "WIRApplicationDictionaryKey"
  @wir_application_identifier_key "WIRApplicationIdentifierKey"
  @wir_application_bundle_identifier_key "WIRApplicationBundleIdentifierKey"
  @wir_is_application_ready_key "WIRIsApplicationReadyKey"
  @wir_session_identifier_key "WIRSessionIdentifierKey"
  @wir_session_capabilities_key "WIRSessionCapabilitiesKey"
  @wir_listing_key "WIRListingKey"
  @wir_type_key "WIRTypeKey"
  @wir_type_automation "WIRTypeAutomation"
  @wir_page_identifier_key "WIRPageIdentifierKey"
  @wir_sender_key "WIRSenderKey"
  @wir_socket_data_key "WIRSocketDataKey"
  @wir_destination_key "WIRDestinationKey"
  @wir_message_data_key "WIRMessageDataKey"

  typedstruct do
    @typedoc false
    # The socket connected to the web-inspector service
    field :ssl_sock, :ssl.sslsocket()
    # The unique session id identifying this connection
    field :session_id, String.t()
    # The pid of the controlling process
    field :cp_pid, reference()
    # The monitor reference to the controlling process
    field :cp_mref, reference()
    # The reported state of the safari process, if any
    field :safari, %{String.t() => any} | nil
    # The id of the currently 'automatable' page.
    field :page_id, String.t() | nil
    # The sequence number of the next request to the page
    field :page_out, non_neg_integer(), default: 0
    # A map of pending replies to page requests. The key is the
    # sequence number and the value is an arbitrary context.
    field :page_in, %{non_neg_integer() => any()}, default: %{}
    # Fail initialization if safari is already running
    field :fail_if_safari_running, boolean, required: true
  end

  @impl true
  def callback_mode, do: :handle_event_function

  @impl true
  def init(args) do
    {:ok, :created, nil, {:next_event, :internal, {:init, args}}}
  end

  @impl true

  #
  # Handle api calls
  #
  def handle_event({:call, {caller, _} = from}, {:set_cp, cp}, _, %__MODULE__{} = data) do
    if caller == data.cp_pid do
      # The controlling process can transfer control to another process
      mref = Process.monitor(cp)
      Process.demonitor(data.cp_mref, [:flush])
      {:keep_state, %__MODULE__{data | cp_pid: cp, cp_mref: mref}, {:reply, from, :ok}}
    else
      {:keep_state_and_data, {:reply, from, {:error, :not_controlling_process}}}
    end
  end

  def handle_event(
        {:call, from},
        {:wait_for_session, _},
        :failed,
        %__MODULE__{fail_if_safari_running: true} = data
      )
      when data.safari != nil do
    {:keep_state_and_data, {:reply, from, {:error, :running}}}
  end

  def handle_event({:call, from}, _, :failed, %__MODULE__{}) do
    # If the process failed to initialize, all other requests will fail
    {:keep_state_and_data, {:reply, from, {:error, :failed}}}
  end

  def handle_event({:call, _}, _, state, %__MODULE__{}) when state != :connected do
    # All requests must otherwise wait for the process to connect to safari
    {:keep_state_and_data, :postpone}
  end

  def handle_event({:call, from}, :create_page, _, %__MODULE__{} = data) do
    send_rpc("createBrowsingContext", [], from, data, fn
      %{"result" => %{"handle" => handle}}, data ->
        {:keep_state, data, {:reply, from, {:ok, handle}}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, :list_pages, _, %__MODULE__{} = data) do
    send_rpc("getBrowsingContexts", [], from, data, fn
      %{"result" => %{"contexts" => contexts}}, data ->
        pages =
          for %{"active" => active, "handle" => handle, "url" => url} <- contexts do
            %{active: active, id: handle, url: url}
          end

        {:keep_state, data, {:reply, from, {:ok, pages}}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:navigate_to, page, url, opts}, _, %__MODULE__{} = data) do
    args = [handle: page, url: url] ++ page_load_timeout(opts)

    send_rpc("navigateBrowsingContext", args, from, data, fn
      %{"result" => %{}}, data ->
        {:keep_state, data, {:reply, from, :ok}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:switch_to_page, page}, _, %__MODULE__{} = data) do
    args = [browsingContextHandle: page, frameHandle: ""]

    send_rpc("switchToBrowsingContext", args, from, data, fn
      %{"result" => %{}}, data ->
        {:keep_state, data, {:reply, from, :ok}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:take_screenshot, page}, _, %__MODULE__{} = data) do
    args = [handle: page, scrollIntoViewIfNeeded: true, clipToViewport: true]

    send_rpc("takeScreenshot", args, from, data, fn
      %{"result" => %{"data" => base64}}, data ->
        {:keep_state, data, {:reply, from, {:ok, Base.decode64!(base64)}}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:go_forward, page}, _, %__MODULE__{} = data) do
    send_rpc("goForwardInBrowsingContext", [handle: page], from, data, fn
      %{"result" => %{}}, data ->
        {:keep_state, data, {:reply, from, :ok}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:go_back, page}, _, %__MODULE__{} = data) do
    send_rpc("goBackInBrowsingContext", [handle: page], from, data, fn
      %{"result" => %{}}, data ->
        {:keep_state, data, {:reply, from, :ok}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:reload, page}, _, %__MODULE__{} = data) do
    send_rpc("reloadBrowsingContext", [handle: page], from, data, fn
      %{"result" => %{}}, data ->
        {:keep_state, data, {:reply, from, :ok}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:close_page, page}, _, %__MODULE__{} = data) do
    send_rpc("closeBrowsingContext", [handle: page], from, data, fn
      %{"result" => %{}}, data ->
        {:keep_state, data, {:reply, from, :ok}}

      %{"error" => error}, data ->
        {:keep_state, data, {:reply, from, {:error, error}}}
    end)
  end

  def handle_event({:call, from}, {:wait_for_session, _}, :connected, _) do
    {:keep_state_and_data, {:reply, from, :ok}}
  end

  def handle_event({:call, _from}, {:wait_for_session, :infinity}, _, _) do
    {:keep_state_and_data, :postpone}
  end

  def handle_event({:call, from}, {:wait_for_session, deadline}, _, _) do
    next_events = [
      :postpone,
      {{:timeout, {:wait_for_session, from}}, deadline, nil, [abs: true]}
    ]

    {:keep_state_and_data, next_events}
  end

  #
  # Event handling
  #
  def handle_event(:internal, {:init, args}, :created, nil) do
    proc = Keyword.fetch!(args, :controlling_process)
    udid = Keyword.fetch!(args, :udid)
    timeout = Keyword.get(args, :timeout, 30_000)
    fail_if_safari_running = Keyword.get(args, :fail_if_safari_running, false)

    case ExMobileDevice.Services.connect(udid, @service, timeout) do
      {:ok, ssl_sock} ->
        session_id = UUID.uuid4() |> String.upcase()

        if automation_enabled?(ssl_sock, session_id) do
          :ok = :ssl.setopts(ssl_sock, active: :once)
          mref = Process.monitor(proc)

          next_events = [
            {{:timeout, :start_session}, timeout, nil},
            {:next_event, :internal, :start_session}
          ]

          {:keep_state,
           %__MODULE__{
             ssl_sock: ssl_sock,
             session_id: session_id,
             cp_pid: proc,
             cp_mref: mref,
             fail_if_safari_running: fail_if_safari_running
           }, next_events}
        else
          {:stop, :no_automation}
        end

      {:error, :timeout} ->
        {:stop, :timeout}

      error ->
        {:stop, {:shutdown, error}}
    end
  end

  def handle_event(:internal, :start_session, state, %__MODULE__{} = data) do
    case state do
      :created ->
        {:keep_state_and_data, :postpone}

      :initialized ->
        start_session_initialized(data)

      :ready ->
        next_event = [{:next_event, :internal, :start_automation_session}]
        {:keep_state_and_data, next_event}
    end
  end

  def handle_event(:internal, :start_automation_session, _, %__MODULE__{} = data) do
    if data.safari[@wir_automation_availability_key] == @wir_automation_available do
      params = %{
        @wir_session_identifier_key => data.session_id,
        @wir_application_identifier_key => data.safari[@wir_application_identifier_key],
        @wir_session_capabilities_key => %{
          "org.webkit.webdriver.webrtc.allow-insecure-media-capture" => true,
          "org.webkit.webdriver.webrtc.suppress-ice-candidate-filtering" => false
        }
      }

      case send_msg(
             data.ssl_sock,
             data.session_id,
             @rpc_forward_automation_session_request,
             params
           ) do
        :ok ->
          :keep_state_and_data

        _error ->
          {:next_state, :failed, data}
      end
    else
      {:next_state, :failed, data}
    end
  end

  def handle_event(:internal, {:recv, msg}, state, %__MODULE__{} = data) do
    args = msg[@argument]

    case msg[@selector] do
      @rpc_report_connected_applications when state == :created ->
        update = find_safari(Map.values(args[@wir_application_dictionary_key]))
        {:next_state, :initialized, %__MODULE__{data | safari: update}}

      @rpc_application_sent_data ->
        handle_application_sent_data(args, data)

      @rpc_application_sent_listing ->
        app_id = get_in(data.safari, [@wir_application_identifier_key])

        if app_id == args[@wir_application_identifier_key] do
          find_automation_page(args[@wir_listing_key], data)
        else
          :keep_state_and_data
        end

      @rpc_application_disconnected ->
        app_id = get_in(data.safari, [@wir_application_identifier_key])

        if app_id == args[@wir_application_identifier_key] do
          {:keep_state, %__MODULE__{data | safari: nil, page_id: nil}}
        else
          :keep_state_and_data
        end

      rpc when rpc in [@rpc_application_connected, @rpc_application_updated] ->
        handle_application_update(args, state, data)

      _ ->
        :keep_state_and_data
    end
  end

  def handle_event(:internal, :connect_page, _, %__MODULE__{} = data) do
    params = %{
      @wir_sender_key => data.session_id,
      @wir_application_identifier_key => data.safari[@wir_application_identifier_key],
      @wir_page_identifier_key => data.page_id
    }

    case send_msg(data.ssl_sock, data.session_id, @rpc_forward_socket_setup, params) do
      :ok -> :keep_state_and_data
    end
  end

  #
  # Socket handling
  #

  def handle_event(:info, {:ssl, socket, data}, _, %__MODULE__{ssl_sock: socket}) do
    :ok = :ssl.setopts(socket, active: :once)
    {:keep_state_and_data, {:next_event, :internal, {:recv, ExMobileDevice.Plist.decode(data)}}}
  end

  def handle_event(:info, {:ssl_closed, socket}, _, %__MODULE__{ssl_sock: socket}) do
    {:stop, :shutdown}
  end

  def handle_event(:info, {:DOWN, ref, :process, _, _}, _, %__MODULE__{cp_mref: ref}) do
    {:stop, :shutdown}
  end

  def handle_event({:timeout, {:wait_for_session, from}}, _, _, _) do
    {:keep_state_and_data, {:reply, from, {:error, :timeout}}}
  end

  #
  # Session creation timeout
  #

  def handle_event({:timeout, :start_session}, _, _, %__MODULE__{} = data) do
    {:next_state, :failed, data}
  end

  defp start_session_initialized(%__MODULE__{} = data) do
    cond do
      data.fail_if_safari_running && data.safari ->
        {:next_state, :failed, data}

      get_in(data.safari, [@wir_automation_availability_key]) == @wir_automation_available ->
        {:next_state, :ready, data, :postpone}

      true ->
        launch_safari(data)
    end
  end

  defp launch_safari(%__MODULE__{} = data) do
    bundle_key = %{@wir_application_bundle_identifier_key => @bundle}

    case send_msg(data.ssl_sock, data.session_id, @rpc_request_application_launch, bundle_key) do
      :ok -> {:keep_state_and_data, :postpone}
      _error -> {:next_state, :failed, data}
    end
  end

  defp handle_application_sent_data(args, %__MODULE__{} = data) do
    if args[@wir_destination_key] != data.session_id do
      :keep_state_and_data
    else
      case Jason.decode!(args[@wir_message_data_key]) do
        %{"id" => id} = response ->
          dispatch_page_response(id, response, data)

        _ ->
          :keep_state_and_data
      end
    end
  end

  defp dispatch_page_response(id, response, %__MODULE__{} = data) do
    case Map.pop(data.page_in, id) do
      {nil, _} ->
        :keep_state_and_data

      {fun, pending} ->
        fun.(response, %__MODULE__{data | page_in: pending})
    end
  end

  defp handle_application_update(args, state, %__MODULE__{} = data) do
    if args[@wir_application_bundle_identifier_key] != @bundle do
      :keep_state_and_data
    else
      if state == :initialized && safari_automated?(args) && safari_ready?(args) do
        {:next_state, :ready, %__MODULE__{data | safari: args}}
      else
        {:keep_state, %__MODULE__{data | safari: args}}
      end
    end
  end

  defp find_safari(applications) do
    Enum.find(applications, &match?(%{@wir_application_bundle_identifier_key => @bundle}, &1))
  end

  defp safari_automated?(%{@wir_automation_availability_key => @wir_automation_available}),
    do: true

  defp safari_automated?(_), do: false

  defp safari_ready?(safari), do: safari[@wir_is_application_ready_key]

  defp send_rpc(method, args, from, %__MODULE__{} = data, fun) do
    seqno = data.page_out

    case send_cmd(data, seqno, method, args) do
      :ok ->
        page_in = Map.put(data.page_in, seqno, fun)
        {:keep_state, %{data | page_out: seqno + 1, page_in: page_in}}

      error ->
        stop_and_reply(error, from)
    end
  end

  defp send_cmd(%__MODULE__{} = data, id, method, args) do
    call_args = %{
      "method" => "Automation.#{method}",
      "params" => Map.new(args),
      "id" => id
    }

    call_params = %{
      @wir_application_identifier_key => data.safari[@wir_application_identifier_key],
      @wir_page_identifier_key => data.page_id,
      @wir_session_identifier_key => data.session_id,
      @wir_socket_data_key => {:data, Jason.encode!(call_args)}
    }

    send_msg(data.ssl_sock, data.session_id, @rpc_forward_socket_data, call_params)
  end

  defp find_automation_page(pages, %__MODULE__{session_id: sid} = data) do
    candidates =
      for {_,
           %{@wir_type_key => @wir_type_automation, @wir_session_identifier_key => ^sid} = page} <-
            pages do
        page[@wir_page_identifier_key]
      end

    if pid = List.first(candidates) do
      cond do
        is_nil(data.page_id) ->
          {:keep_state, %__MODULE__{data | page_id: pid}, {:next_event, :internal, :connect_page}}

        get_in(pages, [to_string(data.page_id), @wir_connection_identifier_key]) == sid ->
          {:next_state, :connected, data, {{:timeout, :start_session}, :cancel}}
      end
    else
      :keep_state_and_data
    end
  end

  defp page_load_timeout(opts) do
    if t = Keyword.get(opts, :timeout) do
      [pageLoadTimeout: t]
    else
      []
    end
  end

  defp automation_enabled?(ssl_sock, session_id) do
    with :ok <- send_msg(ssl_sock, session_id, @rpc_report_identifier),
         {:ok, %{@selector => @rpc_report_current_state} = reply} <- recv_msg(ssl_sock, 5000) do
      get_in(reply, [@argument, @wir_automation_availability_key]) == @wir_automation_available
    else
      _ -> false
    end
  end

  defp recv_msg(ssl_sock, timeout) do
    with {:ok, packet} <- :ssl.recv(ssl_sock, 0, timeout) do
      {:ok, ExMobileDevice.Plist.decode(packet)}
    end
  end

  defp send_msg(ssl_sock, conn_id, selector, args \\ %{}) do
    send_plist(ssl_sock, %{
      @selector => selector,
      @argument => Map.put(args, @wir_connection_identifier_key, conn_id)
    })
  end

  defp stop_and_reply(error, to) do
    if to do
      {:stop_and_reply, {:shutdown, error}, {:reply, to, error}}
    else
      {:stop, {:shutdown, error}}
    end
  end

  defp send_plist(ssl_sock, plist) do
    :ssl.send(ssl_sock, ExMobileDevice.Plist.encode(plist))
  end
end