teac

defmodule Teac.WssClient do
  use WebSockex
  require Logger

  @wss_endpoint "wss://eventsub.wss.twitch.tv/ws"

  # Extra buffer added on top of Twitch's keepalive_timeout_seconds before
  # we consider the connection dead.
  @keepalive_buffer_ms 1_000

  @close_codes %{
    4000 => "internal server error",
    4001 => "client sent inbound traffic (prohibited)",
    4002 => "client failed ping-pong",
    4003 => "connection unused — no subscription within 10 seconds",
    4004 => "reconnect grace period expired",
    4005 => "network timeout",
    4006 => "network error",
    4007 => "invalid reconnect URL"
  }

  @doc """
  Returns a child spec for use in a supervision tree.

  The argument must be a two-element tuple of `{client, opts}`, where `opts`
  is the same keyword list accepted by `start_link/2`. A `:name` opt is
  recommended so the process can be looked up by name.

  ## Example

      children = [
        {Teac.WssClient, {client, handler: MyApp.TwitchHandler, name: MyApp.TwitchSocket}}
      ]

      Supervisor.start_link(children, strategy: :one_for_one)
  """
  def child_spec({client, opts}) do
    %{
      id: Keyword.get(opts, :name, __MODULE__),
      start: {__MODULE__, :start_link, [client, opts]},
      restart: Keyword.get(opts, :restart, :permanent),
      type: :worker
    }
  end

  @doc """
  Starts a WebSocket connection to the Twitch EventSub endpoint.

  ## Options

    * `:handler` - required. A module implementing `Teac.WssClient.Handler`.
    * `:handler_state` - optional initial state passed to handler callbacks. Defaults to `nil`.
    * `:keepalive_timeout_seconds` - optional integer (10–600). Tells Twitch how often to
      send keepalive frames. Defaults to Twitch's server-side default (10s).

  ## Example

      Teac.WssClient.start_link(client,
        handler: MyApp.TwitchHandler,
        keepalive_timeout_seconds: 30
      )
  """
  def start_link(%Teac.Client{} = client, opts \\ []) do
    handler = Keyword.fetch!(opts, :handler)

    url =
      case Keyword.get(opts, :url) do
        nil -> build_url(Keyword.get(opts, :keepalive_timeout_seconds))
        url -> url
      end

    websockex_opts =
      Keyword.drop(opts, [
        :url,
        :notify_on_welcome,
        :reconnect_handler,
        :handler,
        :handler_state,
        :keepalive_timeout_seconds
      ])

    initial_state = %{
      client: client,
      session_id: nil,
      notify_on_welcome: Keyword.get(opts, :notify_on_welcome),
      reconnect_handler: Keyword.get(opts, :reconnect_handler),
      handler: handler,
      handler_state: Keyword.get(opts, :handler_state),
      keepalive_timer: nil,
      keepalive_timeout_ms: nil,
      seen_message_ids: MapSet.new(),
      reconnect_attempts: Keyword.get(opts, :reconnect_attempts, 0)
    }

    WebSockex.start_link(url, __MODULE__, initial_state, websockex_opts)
  end

  @doc """
  Gracefully closes the WebSocket connection with a normal close frame.
  """
  def stop(pid) do
    GenServer.call(pid, :stop)
  end

  @doc """
  Creates an EventSub subscription tied to this WebSocket session.

  Must be called after the connection is established (i.e. after `session_welcome`
  has been received). Twitch requires at least one subscription within 10 seconds
  of the welcome message or it closes the connection with code 4003.

  ## Options

    * `:type` - required. The subscription type, e.g. `"channel.follow"`.
    * `:version` - required. The subscription version, e.g. `"2"`.
    * `:condition` - required. A map of condition parameters for the subscription type.

  ## Example

      Teac.WssClient.subscribe(pid,
        type: "channel.follow",
        version: "2",
        condition: %{"broadcaster_user_id" => "123", "moderator_user_id" => "123"}
      )
  """
  def subscribe(pid, opts) do
    GenServer.call(pid, {:subscribe, opts})
  end

  @doc """
  Assigns this WebSocket session to a conduit shard.

  For conduit-based setups, call this instead of `subscribe/2` within 10 seconds
  of `session_welcome`. Subscriptions are then created separately via
  `Teac.Api.EventSub.Subscriptions.post/2` using the conduit ID as transport.

  ## Options

    * `:conduit_id` - required. The conduit to assign this session to.
    * `:shard_id` - required. The shard ID (zero-based string) to assign.

  ## Example

      Teac.WssClient.assign_shard(pid, conduit_id: "abc123", shard_id: "0")
  """
  def assign_shard(pid, opts) do
    GenServer.call(pid, {:assign_shard, opts})
  end

  @doc """
  Deletes an EventSub subscription by ID.

  ## Example

      Teac.WssClient.unsubscribe(pid, "2d9e9f1b-9b7a-4e6a-8f3c-1a2b3c4d5e6f")
  """
  def unsubscribe(pid, subscription_id) do
    GenServer.call(pid, {:unsubscribe, subscription_id})
  end

  @impl true
  def handle_connect(_conn, state) do
    Logger.info("Connected to Twitch EventSub WebSocket.")
    {:ok, state}
  end

  @impl true
  def handle_disconnect(%{code: 1000}, state) do
    Logger.info("WebSocket closed normally.")
    {:ok, state}
  end

  @impl true
  def handle_disconnect(%{code: 1011}, state) do
    Logger.info("WebSocket closed due to keepalive timeout — reconnecting.")
    {:ok, state}
  end

  @impl true
  def handle_disconnect(%{code: code} = _status, _state) when is_map_key(@close_codes, code) do
    reason = Map.fetch!(@close_codes, code)
    Logger.error("Twitch closed connection (#{code}): #{reason}")
    exit({:twitch_close, code, reason})
  end

  @impl true
  def handle_disconnect(status, state) do
    Logger.warning("WebSocket disconnected: #{inspect(status)}")
    {:ok, state}
  end

  @impl true
  def handle_frame({_type, msg}, state) do
    case JSON.decode(msg) do
      {:ok, data} ->
        handle_message(data, state)

      {:error, reason} ->
        Logger.error("Failed to decode JSON: #{inspect(reason)}")
        {:ok, state}
    end
  end

  # WebSockex routes GenServer.call/3 through handle_info as {:"$gen_call", from, msg}.
  @impl true
  def handle_info({:"$gen_call", from, {:subscribe, _opts}}, %{session_id: nil} = state) do
    GenServer.reply(from, {:error, "WebSocket session not yet established"})
    {:ok, state}
  end

  def handle_info({:"$gen_call", from, {:subscribe, opts}}, state) do
    payload = %{
      "type" => Keyword.fetch!(opts, :type),
      "version" => Keyword.fetch!(opts, :version),
      "condition" => Keyword.fetch!(opts, :condition),
      "transport" => %{"method" => "websocket", "session_id" => state.session_id}
    }

    result = Teac.Api.EventSub.Subscriptions.post(state.client, payload: payload)
    GenServer.reply(from, result)
    {:ok, state}
  end

  def handle_info({:"$gen_call", from, {:unsubscribe, id}}, state) do
    result = Teac.Api.EventSub.Subscriptions.delete(state.client, id: id)
    GenServer.reply(from, result)
    {:ok, state}
  end

  def handle_info({:"$gen_call", from, {:assign_shard, _opts}}, %{session_id: nil} = state) do
    GenServer.reply(from, {:error, "WebSocket session not yet established"})
    {:ok, state}
  end

  def handle_info({:"$gen_call", from, :stop}, state) do
    GenServer.reply(from, :ok)
    {:close, {1000, "stopped"}, state}
  end

  def handle_info({:"$gen_call", from, {:assign_shard, opts}}, state) do
    shard = %{
      "id" => Keyword.fetch!(opts, :shard_id),
      "transport" => %{"method" => "websocket", "session_id" => state.session_id}
    }

    result =
      Teac.Api.EventSub.Conduits.Shards.patch(state.client,
        conduit_id: Keyword.fetch!(opts, :conduit_id),
        shards: [shard]
      )

    GenServer.reply(from, result)
    {:ok, state}
  end

  def handle_info({:do_reconnect, url}, state) do
    case start_link(state.client,
           url: url,
           handler: state.handler,
           handler_state: state.handler_state,
           notify_on_welcome: self(),
           reconnect_attempts: state.reconnect_attempts
         ) do
      {:ok, _pid} ->
        {:ok, state}

      {:error, reason} ->
        Logger.error("Reconnect failed: #{inspect(reason)}")
        {:ok, state}
    end
  end

  def handle_info({:twitch_welcome, _new_pid}, state) do
    Logger.info("Reconnect complete, closing old connection.")
    {:close, {1000, "session_reconnect"}, state}
  end

  def handle_info(:keepalive_timeout, state) do
    Logger.error("Keepalive timeout — no message received within the expected window. Closing.")
    {:close, {1011, "keepalive_timeout"}, state}
  end

  defp handle_message(
         %{"metadata" => %{"message_id" => msg_id, "message_type" => type}} = data,
         state
       ) do
    if MapSet.member?(state.seen_message_ids, msg_id) do
      Logger.debug("Dropping duplicate message: #{msg_id}")
      {:ok, state}
    else
      state = %{state | seen_message_ids: MapSet.put(state.seen_message_ids, msg_id)}
      dispatch_message(type, data, state)
    end
  end

  defp handle_message(_data, state), do: {:ok, state}

  defp dispatch_message("session_welcome", data, state) do
    session_id = get_in(data, ["payload", "session", "id"])
    timeout_s = get_in(data, ["payload", "session", "keepalive_timeout_seconds"])

    Logger.info(
      "Session established: #{session_id} (keepalive: #{timeout_s}s). Subscribe within 10 seconds."
    )

    if state.notify_on_welcome do
      send(state.notify_on_welcome, {:twitch_welcome, self()})
    end

    ms = timeout_s * 1_000 + @keepalive_buffer_ms

    state =
      state
      |> Map.put(:session_id, session_id)
      |> Map.put(:notify_on_welcome, nil)
      |> Map.put(:keepalive_timeout_ms, ms)
      |> schedule_keepalive(ms)

    {:ok, state}
  end

  defp dispatch_message("session_keepalive", _data, state) do
    {:ok, reset_keepalive(state)}
  end

  defp dispatch_message("notification", data, state) do
    subscription = get_in(data, ["payload", "subscription"])
    event = get_in(data, ["payload", "event"])
    Logger.debug("Received notification: #{subscription["type"]}")

    state = reset_keepalive(state)

    case state.handler.handle_event(subscription, event, state.handler_state) do
      {:ok, new_handler_state} ->
        {:ok, %{state | handler_state: new_handler_state}}

      {:error, reason} ->
        Logger.error("Handler error for #{subscription["type"]}: #{inspect(reason)}")
        {:ok, state}
    end
  end

  defp dispatch_message("session_reconnect", data, %{reconnect_handler: handler} = state)
       when not is_nil(handler) do
    url = get_in(data, ["payload", "session", "reconnect_url"])
    Logger.info("Reconnect requested, forwarding to reconnect_handler.")
    send(handler, {:twitch_reconnect, self(), url})
    {:ok, state}
  end

  defp dispatch_message("session_reconnect", data, state) do
    url = get_in(data, ["payload", "session", "reconnect_url"])
    delay = backoff_ms(state.reconnect_attempts)
    Logger.info("Reconnect requested, connecting in #{delay}ms: #{url}")
    Process.send_after(self(), {:do_reconnect, url}, delay)
    {:ok, %{state | reconnect_attempts: state.reconnect_attempts + 1}}
  end

  defp dispatch_message("revocation", data, state) do
    subscription = get_in(data, ["payload", "subscription"])

    Logger.warning(
      "Subscription revoked — type: #{subscription["type"]}, reason: #{subscription["status"]}"
    )

    state.handler.handle_revocation(subscription, state.handler_state)
    {:ok, state}
  end

  defp dispatch_message(other, _data, state) do
    Logger.warning("Unhandled message type: #{other}")
    {:ok, state}
  end

  # Exponential backoff: 1s, 2s, 4s, 8s, capped at 30s.
  defp backoff_ms(attempts), do: min(:math.pow(2, attempts) * 1_000, 30_000) |> round()

  defp build_url(nil), do: @wss_endpoint
  defp build_url(seconds), do: "#{@wss_endpoint}?keepalive_timeout_seconds=#{seconds}"

  defp schedule_keepalive(state, ms) do
    if state.keepalive_timer, do: Process.cancel_timer(state.keepalive_timer)
    %{state | keepalive_timer: Process.send_after(self(), :keepalive_timeout, ms)}
  end

  defp reset_keepalive(%{keepalive_timeout_ms: nil} = state), do: state

  defp reset_keepalive(state) do
    Process.cancel_timer(state.keepalive_timer)

    %{
      state
      | keepalive_timer:
          Process.send_after(self(), :keepalive_timeout, state.keepalive_timeout_ms)
    }
  end
end