polymarket_sdk

defmodule Polymarket.MarketStream do
  @moduledoc """
  CLOB **market-data** websocket client — the order-book stream.

  Connects to `wss://ws-subscriptions-clob.polymarket.com/ws/market`,
  subscribes to a set of `asset_ids` (outcome token ids), and forwards
  book updates to a callback PID:

    * `{:polymarket_market, :book, payload}` — a full book snapshot
      (`%{"asset_id", "bids", "asks", ...}`), sent on subscribe and
      periodically;
    * `{:polymarket_market, :price_change, payload}` — incremental level
      changes (`%{"asset_id", "changes": [%{"price","side","size"}]}`).

  This is the push-based alternative to polling `GET /book` — one
  persistent connection feeds book updates for all subscribed tokens,
  so the snapshot layer never has to poll (which can be rate-limiting).
  Book *maintenance* (applying incremental `:price_change` events to a
  stored book snapshot) is the consumer's job.

  Mirrors `Polymarket.RTDS` (WebSockex, PID callback, text PING) plus a
  **reconnect backoff** so transient drops / `429`s don't storm.

  ## Options

    * `:callback` *(required)* — PID to receive events.
    * `:assets` *(required, non-empty)* — token ids to subscribe.
    * `:url` — override (default the CLOB market channel).
    * `:ping_interval` — ms (default 10_000).
    * `:name` — optional process name.
  """

  use WebSockex

  require Logger

  @url "wss://ws-subscriptions-clob.polymarket.com/ws/market"
  @default_ping_interval_ms 10_000
  @default_resubscribe_interval_ms 90_000
  @max_backoff_ms 30_000

  @spec start_link(keyword()) :: {:ok, pid()} | {:error, term()}
  def start_link(opts) do
    callback = validate_callback!(opts)
    assets = validate_assets!(opts)
    url = Keyword.get(opts, :url, @url)
    ping = Keyword.get(opts, :ping_interval, @default_ping_interval_ms)

    state = %{
      callback: callback,
      assets: assets,
      ping_interval_ms: ping,
      resubscribe_interval_ms:
        Keyword.get(opts, :resubscribe_interval, @default_resubscribe_interval_ms),
      ping_timer: nil,
      resub_timer: nil
    }

    ws_opts = if name = Keyword.get(opts, :name), do: [name: name], else: []
    WebSockex.start_link(url, __MODULE__, state, ws_opts)
  end

  @spec default_url() :: String.t()
  def default_url, do: @url

  @doc """
  Decode a raw CLOB market frame into a list of `{event_type, payload}`
  tuples. Pure — the frame can be a single event object or a JSON array
  of them. Unknown/`PONG`/non-book frames yield `[]`.
  """
  @spec decode_events(String.t()) :: [{atom(), map()}]
  def decode_events(text) when is_binary(text) do
    case Jason.decode(text) do
      {:ok, list} when is_list(list) -> Enum.flat_map(list, &event/1)
      {:ok, %{} = obj} -> event(obj)
      _ -> []
    end
  end

  defp event(%{"event_type" => "book"} = e), do: [{:book, e}]
  defp event(%{"event_type" => "price_change"} = e), do: [{:price_change, e}]
  defp event(_other), do: []

  # ── WebSockex ──

  @impl true
  def handle_connect(_conn, state) do
    Logger.info("[Polymarket.MarketStream] Connected (#{length(state.assets)} assets)")
    # Cancel any timers left over from a previous connection before
    # scheduling fresh ones — otherwise every reconnect leaks a ping and
    # a resubscribe timer, and the frames pile up after a few flaps.
    state = state |> reschedule_ping() |> reschedule_resubscribe()
    Process.send_after(self(), :subscribe, 0)
    {:ok, state}
  end

  @impl true
  def handle_info(:subscribe, state) do
    {:reply, {:text, subscribe_msg(state)}, state}
  end

  # Periodically re-send the subscription so the venue re-pushes fresh
  # book snapshots — bounds staleness even if incremental updates are
  # sparse. Far cheaper than HTTP polling (one frame for all assets).
  def handle_info(:resubscribe, state) do
    state = reschedule_resubscribe(state)
    {:reply, {:text, subscribe_msg(state)}, state}
  end

  def handle_info(:ping, state) do
    state = reschedule_ping(state)
    {:reply, {:text, "PING"}, state}
  end

  def handle_info(_msg, state), do: {:ok, state}

  @impl true
  def handle_frame({:text, "PONG"}, state), do: {:ok, state}

  def handle_frame({:text, msg}, state) do
    for {type, payload} <- decode_events(msg) do
      send(state.callback, {:polymarket_market, type, payload})
    end

    {:ok, state}
  end

  def handle_frame(_frame, state), do: {:ok, state}

  @impl true
  def handle_disconnect(status, state) do
    reason = Map.get(status, :reason)
    attempt = Map.get(status, :attempt_number, 1)
    delay = backoff_ms(attempt)

    # Cancel the ping/resubscribe timers scheduled against the now-dead
    # connection so their refs don't leak; `handle_connect/2` schedules
    # fresh ones on the next successful connect.
    state = state |> cancel_ping() |> cancel_resubscribe()

    Logger.warning(
      "[Polymarket.MarketStream] Disconnected (#{inspect(reason)}); " <>
        "reconnect in #{delay}ms (attempt #{attempt})"
    )

    # Exponential backoff between attempts (first attempt immediate), so
    # a hard-down endpoint / 429 storm doesn't hot-loop. WebSockex's
    # `attempt_number` resets to 1 on a successful connect, so a single
    # transient drop reconnects immediately rather than sleeping.
    if delay > 0, do: Process.sleep(delay)

    {:reconnect, state}
  end

  # ── Private ──

  defp validate_callback!(opts) do
    case Keyword.get(opts, :callback) do
      pid when is_pid(pid) -> pid
      other -> raise ArgumentError, "MarketStream :callback must be a PID, got #{inspect(other)}"
    end
  end

  defp validate_assets!(opts) do
    case Keyword.get(opts, :assets) do
      [_ | _] = assets -> Enum.map(assets, &to_string/1)
      other -> raise ArgumentError, "MarketStream :assets must be a non-empty list, got #{inspect(other)}"
    end
  end

  defp subscribe_msg(state), do: Jason.encode!(%{"assets_ids" => state.assets, "type" => "market"})

  # Cancel-then-reschedule the ping timer, tracking the ref in state so
  # reconnects can't leak overlapping ping loops.
  defp reschedule_ping(%{ping_interval_ms: ms} = state) when is_integer(ms) and ms > 0 do
    state = cancel_ping(state)
    %{state | ping_timer: Process.send_after(self(), :ping, ms)}
  end

  defp reschedule_ping(state), do: cancel_ping(state)

  defp cancel_ping(state) do
    case Map.get(state, :ping_timer) do
      ref when is_reference(ref) -> Process.cancel_timer(ref)
      _ -> :ok
    end

    %{state | ping_timer: nil}
  end

  defp reschedule_resubscribe(%{resubscribe_interval_ms: ms} = state)
       when is_integer(ms) and ms > 0 do
    state = cancel_resubscribe(state)
    %{state | resub_timer: Process.send_after(self(), :resubscribe, ms)}
  end

  defp reschedule_resubscribe(state), do: cancel_resubscribe(state)

  defp cancel_resubscribe(state) do
    case Map.get(state, :resub_timer) do
      ref when is_reference(ref) -> Process.cancel_timer(ref)
      _ -> :ok
    end

    %{state | resub_timer: nil}
  end

  # First reconnect attempt is immediate (0ms); afterwards back off
  # exponentially: 1s, 2s, 4s, … capped at @max_backoff_ms.
  defp backoff_ms(attempt) when is_integer(attempt) and attempt <= 1, do: 0

  defp backoff_ms(attempt) when is_integer(attempt) do
    min(@max_backoff_ms, 1_000 * Bitwise.bsl(1, attempt - 2))
  end
end