sway_sock

defmodule SwaySock do
  @type json :: :json.decode_value()

  @moduledoc """
  SwaySock is a library for controlling the running [Sway](https://swaywm.org/) process through its IPC interface.

  You can query the state of the window manager, listen to events, and execute sway commands. 
  This module provides the main API to interfacing with Sway via IPC

  ## Overview

  `start_link/1` starts a Supervisor that connects to Sway's Unix Domain Socket. This implies the library is useless without a currently running sway instace.

  The architecture is simple. All commands apart from `subscribe/3` are translated to IPC messages and sent to the sway socket.
  The response is immediately read from the socket and returned as-is.

  `subscribe/3` is a special case. The [`i3` docs](https://i3wm.org/docs/ipc.html#_events) explain it clearly:

  > Caveat: As soon as you subscribe to an event, it is not guaranteed any longer that the requests to i3 are processed in order.
  > This means, the following situation can happen: You send a GET_WORKSPACES request but you receive a "workspace" event before receiving the reply to GET_WORKSPACES.
  > If your program does not want to cope which such kinds of race conditions (an event based library may not have a problem here), I suggest you create a separate connection to receive events.

  The same applies to Sway. This library opts to open a new connection to the sway socket for each event subscription. This is always done in supervised Task.

  ## Reply structure

  The sway developers may change the format of replies across releases.
  Users of this library are encouraged to refer to the sway-ipc manual (`man sway-ipc`) for the schemas of the JSON sway sends back.

  ## Examples

  Check out the `examples/` directory. 

  """

  use Supervisor

  @message_types [
    run_command: 0,
    get_workspaces: 1,
    subscribe: 2,
    get_outputs: 3,
    get_tree: 4,
    get_marks: 5,
    get_bar_config: 6,
    get_version: 7,
    get_binding_modes: 8,
    get_config: 9,
    send_tick: 10,
    get_binding_state: 12,
    get_inputs: 100,
    get_seats: 101
  ]

  @event_types [
    workspace: 0x80000000,
    output: 0x80000001,
    mode: 0x80000002,
    window: 0x80000003,
    barconfig_update: 0x80000004,
    binding: 0x80000005,
    shutdown: 0x80000006,
    tick: 0x80000007,
    bar_state_update: 0x80000014,
    input: 0x80000015
  ]

  @impl true
  def init(name) do
    # I start a Task.Supervisor because the `subscribe()` message requires spinning
    # up new Elixir processes, and it would be nice if they were supervised.
    #
    # I start an Agent because I need to store the socket somewhere.
    socket = get_socket()

    children = [
      %{
        id: Agent,
        start: {Agent, :start_link, [fn -> socket end, [name: get_agent_name(name)]]}
      },
      {Task.Supervisor, name: get_task_supervisor_name(name)}
    ]

    Supervisor.init(children, strategy: :one_for_one)
  end

  @spec start_link(name :: atom()) :: {:ok, pid()}
  @doc """
  Connects to the Sway socket. `name` identifies the connection

  Multiple instances of this library may be used as long as they have different `name`s.

  ## How to supervise

      iex> children = [ {SwaySock, :my_socket} ]
      ...> Supervisor.start_link(children, strategy: :one_for_one)

  """
  def start_link(name) when is_atom(name) do
    Supervisor.start_link(__MODULE__, name, name: name)
  end

  @spec run_command(atom(), binary()) :: :ok | :error
  @doc """
  Runs the sway commands in `script` 

  Returns `{:error, <errors>}` if any of the commands in `script` fail. The second element of the tuple will be the list of errors reported by sway.
  Returns `:ok` otherwise.
  """
  def run_command(conn, script) when is_atom(conn) and is_binary(script) do
    results = send_and_receive(conn, :run_command, script)
    succeeded = Enum.map(results, fn result -> result["success"] end) |> Enum.all?()
    errors = Enum.map(results, fn result -> result["error"] || [] end) |> List.flatten()

    if succeeded do
      :ok
    else
      {:error, errors}
    end
  end

  @spec get_workspaces(atom()) :: json()
  @doc """
  Retrieves the list of active workspaces
  """
  def get_workspaces(conn) when is_atom(conn) do
    send_and_receive(conn, :get_workspaces)
  end

  @doc """
  Creates a new linked process that invokes `callback` when any of the listed `event_types` occurs.

  ## Events
  You may subscribe to any of the following events

  `#{@event_types |> Enum.map(fn {t, _} -> ":#{t}" end) |> Enum.join(", ")}`

  Events are described in the sway-ipc manual under the "EVENTS" section.

  ## Callback

  The callback must accept two arguments:

  - The first is the the event details, and has this structure:  `{:event_name, payload}`
    It is a two-element tuple where the first element is an atom representing the event, and the second is its payload. 

  - The second argument is state for the callback. The return value of one invocation of the callback is used as the state for the next.
    The state is initialized to `callback_start_state`. 

  """
  def subscribe(conn, event_types, callback, callback_start_state \\ nil)
      when is_atom(conn) and
             is_list(event_types) and
             is_function(callback, 2) do
    event_types = MapSet.new(event_types)
    allowed = Keyword.keys(get_event_types()) |> MapSet.new()

    unrecognized = MapSet.difference(event_types, allowed)

    if MapSet.size(unrecognized) > 0 do
      raise(
        ArgumentError,
        "The following event types were not recognized #{inspect(unrecognized)}"
      )
    end

    pid = self()
    task_supervisor = get_task_supervisor_name(conn)

    Task.Supervisor.start_child(task_supervisor, fn ->
      setup_subscription(pid, event_types, callback, callback_start_state)
    end)

    receive do
      %{"success" => true} -> :ok
      %{"success" => false} -> :error
    end
  end

  # Creates the new socket, subscribes it to `events`, and begins the subscriber loop
  # Reports the result of subscribing back to the parent process.
  defp setup_subscription(parent_pid, event_types, callback, start_state) do
    socket = get_socket()
    type_id = get_message_types()[:subscribe]

    # Note: event_types is a MapSet, which is an Elixir struct.
    # Make sure to convert it to a list before encoding as JSON
    send_message(socket, type_id, :json.encode(event_types |> MapSet.to_list()))

    case recv_message(socket, type_id) do
      %{"success" => true} = result ->
        send(parent_pid, result)

      response ->
        send(parent_pid, response)
        exit(:normal)
    end

    # We have a map of event types/names -> type ids
    # Construct a reverse map so can take a type id from an incoming event and quickly match that with a human-readable event name
    type_id_lookup = Map.new(get_event_types(), fn {event, type_id} -> {type_id, event} end)
    subscriber_loop(socket, type_id_lookup, event_types, callback, start_state)
  end

  defp subscriber_loop(socket, type_id_lookup, event_types, callback, state) do
    {type_id, payload} = parse_message(socket)
    {:ok, event_type} = Map.fetch(type_id_lookup, type_id)

    if not MapSet.member?(event_types, event_type) do
      raise(
        "Received unexpected event '#{event_type}' in subscription for '#{inspect(event_types)}'"
      )
    end

    state = callback.({event_type, :json.decode(payload)}, state)
    subscriber_loop(socket, type_id_lookup, event_types, callback, state)
  end

  @spec get_outputs(atom()) :: json()
  @doc """
  Returns the list of outputs
  """
  def get_outputs(conn) when is_atom(conn) do
    send_and_receive(conn, :get_outputs)
  end

  @spec get_tree(atom()) :: json()
  @doc """
  Returns the JSON representation of sway's node tree
  """
  def get_tree(conn) when is_atom(conn) do
    send_and_receive(conn, :get_tree)
  end

  @spec get_marks(atom()) :: json()
  @doc """
  Returns the currently set marks
  """
  def get_marks(conn) when is_atom(conn) do
    send_and_receive(conn, :get_marks)
  end

  @spec get_bar_config(atom(), binary()) :: json()
  @doc """
  Returns the list of configured bar IDs.

  When `bar_id` is non-empty, this function returns the configuration of the given bar instead. 
  """
  def get_bar_config(conn, bar_id \\ "") when is_atom(conn) and is_binary(bar_id) do
    send_and_receive(conn, :get_bar_config, bar_id)
  end

  @spec get_version(atom()) :: json()
  @doc """
  Returns version information about the current sway process
  """
  def get_version(conn) when is_atom(conn) do
    send_and_receive(conn, :get_version)
  end

  @spec get_binding_modes(atom()) :: json()
  @doc """
  Returns a list of configured binding modes
  """
  def get_binding_modes(conn) when is_atom(conn) do
    send_and_receive(conn, :get_binding_modes)
  end

  @spec get_config(atom()) :: json()
  @doc """
  Returns the contents of the last-loaded sway configuration
  """
  def get_config(conn) when is_atom(conn) do
    send_and_receive(conn, :get_config)
  end

  @spec send_tick(atom(), binary()) :: :ok | :error
  @doc """
  Sends a TICK event to all clients subscribing to the event to ensure that all events prior to the tick were received.

  If a payload is given, it will be included in the TICK event.
  Returns `:ok` if the message was successfully sent. Returns `:error` otherwise

  """
  def send_tick(conn, payload \\ "") when is_atom(conn) and is_binary(payload) do
    case send_and_receive(conn, :send_tick, payload) do
      %{"success" => true} -> :ok
      %{"success" => false} -> :error
    end
  end

  @spec get_binding_state(atom()) :: json()
  @doc """
  Returns the currently active binding mode.
  """
  def get_binding_state(conn) when is_atom(conn) do
    send_and_receive(conn, :get_binding_state)
  end

  @spec get_inputs(atom()) :: json()
  @doc """
  Returns a list of the input devices currently available
  """
  def get_inputs(conn) when is_atom(conn) do
    send_and_receive(conn, :get_inputs)
  end

  defp get_message_types() do
    @message_types
  end

  defp get_event_types() do
    @event_types
  end

  # Parse an IPC message from the given socket
  defp parse_message(socket) do
    {:ok, header} = :gen_tcp.recv(socket, 14)

    <<"i3-ipc", len::native-size(32), type::native-size(32)>> =
      header

    {:ok, payload} = :gen_tcp.recv(socket, len)

    {type, payload}
  end

  # Recieves a message of a particular type from the sway IPC socket and interprets it as json
  defp recv_message(socket, type_id) when is_port(socket) and is_integer(type_id) do
    {^type_id, payload} = parse_message(socket)
    :json.decode(payload)
  end

  # Sends a message to the sway IPC socket 
  defp send_message(socket, type_id, payload)
       when is_port(socket) and is_list(payload)
       when is_port(socket) and is_binary(payload) do
    allowed = get_message_types() |> Enum.map(fn {_, t} -> t end)

    if type_id not in allowed do
      raise(ArgumentError, "message type '#{type_id}' is not recognized")
    end

    payload_length = IO.iodata_length(payload)
    msg = ["i3-ipc", <<payload_length::native-size(32)>>, <<type_id::native-size(32)>>, payload]
    :ok = :gen_tcp.send(socket, msg)
  end

  defp send_and_receive(conn, type, payload \\ []) when is_atom(conn) and is_atom(type) do
    socket = Agent.get(get_agent_name(conn), fn socket -> socket end)
    type_id = Keyword.fetch!(get_message_types(), type)
    send_message(socket, type_id, payload)
    recv_message(socket, type_id)
  end

  defp get_socket() do
    path = get_socket_path()
    {:ok, socket} = :gen_tcp.connect({:local, path}, 0, [:local, :binary, {:active, false}])
    socket
  end

  defp get_socket_path() do
    case System.fetch_env("SWAYSOCK") do
      {:ok, path} ->
        path

      :error ->
        case System.cmd("sway", ["--get-socketpath"], stderr_to_stdout: true) do
          {path, 0} -> {:ok, path}
          {err, _status} -> raise(SwaySock.NotFound, err)
        end
    end
  end

  defp get_agent_name(conn) do
    String.to_atom("#{conn}__store")
  end

  defp get_task_supervisor_name(conn) do
    String.to_atom("#{conn}__tasks")
  end
end