spacetimedb_ex

defmodule SpacetimeDB.Protocol do
  @moduledoc """
  Encodes and decodes SpacetimeDB WebSocket messages using the
  `v1.json.spacetimedb` text subprotocol.

  ## Wire format

  Both client→server and server→client messages are tagged JSON objects:

      # client → server
      {"CallReducer":   {"reducer": "...", "args": "...", "request_id": N, "flags": 0}}
      {"Subscribe":     {"query_strings": [...], "request_id": N}}
      {"SubscribeSingle": {"query": "...", "request_id": N, "query_id": N}}
      {"SubscribeMulti":  {"query_strings": [...], "request_id": N, "query_id": N}}
      {"Unsubscribe":   {"request_id": N, "query_id": N}}
      {"OneOffQuery":   {"message_id": [...], "query_string": "..."}}

      # server → client
      {"IdentityToken":        {...}}
      {"InitialSubscription":  {...}}
      {"SubscribeApplied":     {...}}
      {"UnsubscribeApplied":   {...}}
      {"SubscriptionError":    {...}}
      {"TransactionUpdate":    {...}}
      {"TransactionUpdateLight": {...}}
      {"OneOffQueryResponse":  {...}}
  """

  alias SpacetimeDB.Types

  @subprotocol "v1.json.spacetimedb"

  @doc "The WebSocket subprotocol header value."
  def subprotocol, do: @subprotocol

  # ---------------------------------------------------------------------------
  # Encoding  (client → server)
  # ---------------------------------------------------------------------------

  @doc """
  Subscribe to multiple queries at once (legacy multi-query subscription).
  The server responds with `InitialSubscription`.
  """
  def encode_subscribe(query_strings, request_id) when is_list(query_strings) do
    Jason.encode!(%{"Subscribe" => %{"query_strings" => query_strings, "request_id" => request_id}})
  end

  @doc """
  Subscribe to a single query. The server responds with `SubscribeApplied`.
  Use `query_id` to identify this subscription in future `Unsubscribe` messages.
  """
  def encode_subscribe_single(query, request_id, query_id) do
    Jason.encode!(%{
      "SubscribeSingle" => %{"query" => query, "request_id" => request_id, "query_id" => query_id}
    })
  end

  @doc """
  Subscribe to multiple queries under a single `query_id`.
  The server responds with `SubscribeMultiApplied`.
  """
  def encode_subscribe_multi(query_strings, request_id, query_id) when is_list(query_strings) do
    Jason.encode!(%{
      "SubscribeMulti" => %{
        "query_strings" => query_strings,
        "request_id" => request_id,
        "query_id" => query_id
      }
    })
  end

  @doc "Remove a subscription previously created with `SubscribeSingle` or `SubscribeMulti`."
  def encode_unsubscribe(request_id, query_id) do
    Jason.encode!(%{"Unsubscribe" => %{"request_id" => request_id, "query_id" => query_id}})
  end

  @doc """
  Call a reducer.

  `args` must be JSON-serialisable (list of reducer arguments in order).
  In the JSON wire format, args are sent as a JSON-encoded string.
  """
  def encode_call_reducer(reducer, args, request_id, flags \\ 0) do
    Jason.encode!(%{
      "CallReducer" => %{
        "reducer" => reducer,
        "args" => Jason.encode!(args),
        "request_id" => request_id,
        "flags" => flags
      }
    })
  end

  @doc """
  Execute a one-off query without establishing a subscription.
  `message_id` is a client-generated 16-byte binary used to correlate the response.
  """
  def encode_one_off_query(query_string, message_id) when is_binary(message_id) do
    Jason.encode!(%{
      "OneOffQuery" => %{
        "message_id" => :binary.bin_to_list(message_id),
        "query_string" => query_string
      }
    })
  end

  # ---------------------------------------------------------------------------
  # Decoding  (server → client)
  # ---------------------------------------------------------------------------

  @doc "Decode a JSON text frame from the server into a typed struct."
  @spec decode(binary()) :: {:ok, term()} | {:error, term()}
  def decode(json) when is_binary(json) do
    case Jason.decode(json) do
      {:ok, map} -> {:ok, decode_map(map)}
      {:error, _} = err -> err
    end
  end

  defp decode_map(%{"IdentityToken" => d}) do
    %Types.IdentityToken{
      identity: d["identity"],
      token: d["token"],
      connection_id: d["connectionId"] || d["connection_id"]
    }
  end

  defp decode_map(%{"InitialSubscription" => d}) do
    %Types.InitialSubscription{
      request_id: d["request_id"],
      tables: decode_table_updates(d["database_update"]["tables"] || []),
      execution_time_micros: d["total_host_execution_time_micros"] || 0
    }
  end

  defp decode_map(%{"SubscribeApplied" => d}) do
    %Types.SubscribeApplied{
      request_id: d["request_id"],
      query_id: get_query_id(d),
      tables: decode_rows_in_tables(d["rows"] || d["tables"] || [])
    }
  end

  defp decode_map(%{"SubscribeMultiApplied" => d}) do
    %Types.SubscribeApplied{
      request_id: d["request_id"],
      query_id: get_query_id(d),
      tables: decode_rows_in_tables(d["rows"] || d["tables"] || [])
    }
  end

  defp decode_map(%{"UnsubscribeApplied" => d}) do
    %Types.UnsubscribeApplied{
      request_id: d["request_id"],
      query_id: get_query_id(d),
      tables: decode_rows_in_tables(d["rows"] || d["tables"] || [])
    }
  end

  defp decode_map(%{"SubscriptionError" => d}) do
    %Types.SubscriptionError{
      request_id: d["request_id"],
      query_id: get_query_id(d),
      error: d["error"]
    }
  end

  defp decode_map(%{"TransactionUpdate" => d}) do
    {status, tables} = decode_update_status(d["status"], d)

    %Types.TransactionUpdate{
      status: status,
      tables: tables,
      timestamp: decode_timestamp(d["timestamp"]),
      caller_identity: d["callerIdentity"] || d["caller_identity"],
      caller_connection_id: d["callerConnectionId"] || d["caller_connection_id"],
      reducer_call: decode_reducer_call(d["reducerCall"] || d["reducer_call"]),
      energy_consumed: d["energyConsumed"] || d["energy_quanta_used"] || 0
    }
  end

  defp decode_map(%{"TransactionUpdateLight" => d}) do
    %Types.TransactionUpdate{
      status: :committed,
      tables: decode_table_updates(d["tables"] || [])
    }
  end

  defp decode_map(%{"OneOffQueryResponse" => d}) do
    %Types.OneOffQueryResponse{
      message_id: decode_message_id(d["messageId"] || d["message_id"]),
      error: d["error"],
      tables: decode_rows_in_tables(d["tables"] || []),
      execution_time_micros: d["total_host_execution_time_micros"] || 0
    }
  end

  defp decode_map(other), do: {:unknown, other}

  # ---------------------------------------------------------------------------
  # Helpers
  # ---------------------------------------------------------------------------

  defp decode_update_status(%{"Committed" => committed}, _d) do
    tables =
      committed["tables"] ||
        get_in(committed, ["database_update", "tables"]) ||
        []

    {:committed, decode_table_updates(tables)}
  end

  defp decode_update_status(%{"Failed" => reason}, _d), do: {{:failed, reason}, []}
  defp decode_update_status(%{"OutOfEnergy" => _}, _d), do: {:out_of_energy, []}
  defp decode_update_status(nil, d), do: {:committed, decode_table_updates(d["tables"] || [])}
  defp decode_update_status(_, _d), do: {:unknown, []}

  defp decode_table_updates(tables) when is_list(tables) do
    Enum.map(tables, fn t ->
      %Types.TableUpdate{
        table_id: t["tableId"] || t["table_id"],
        table_name: t["tableName"] || t["table_name"],
        inserts: decode_rows(t["updates"]["inserts"] || t["inserts"] || []),
        deletes: decode_rows(t["updates"]["deletes"] || t["deletes"] || [])
      }
    end)
  end

  defp decode_table_updates(_), do: []

  # SubscribeApplied / OneOffQueryResponse carry rows differently
  defp decode_rows_in_tables(rows) when is_list(rows) do
    Enum.map(rows, fn
      %{"tableName" => name, "tableRows" => updates} ->
        %Types.TableUpdate{
          table_name: name,
          inserts: decode_rows(updates["inserts"] || []),
          deletes: decode_rows(updates["deletes"] || [])
        }

      %{"table_name" => name} = t ->
        %Types.TableUpdate{
          table_name: name,
          inserts: decode_rows(t["inserts"] || []),
          deletes: decode_rows(t["deletes"] || [])
        }

      other ->
        %Types.TableUpdate{inserts: decode_rows(other["inserts"] || []), deletes: []}
    end)
  end

  defp decode_rows_in_tables(_), do: []

  # Rows in JSON format are JSON-encoded strings of each row value
  defp decode_rows(rows) when is_list(rows) do
    Enum.map(rows, fn
      row when is_binary(row) ->
        case Jason.decode(row) do
          {:ok, decoded} -> decoded
          _ -> row
        end

      row ->
        row
    end)
  end

  defp decode_rows(_), do: []

  defp decode_reducer_call(nil), do: nil

  defp decode_reducer_call(rc) do
    %Types.ReducerCallInfo{
      reducer_name: rc["reducerName"] || rc["reducer_name"],
      request_id: rc["requestId"] || rc["request_id"],
      status: decode_call_status(rc["status"])
    }
  end

  defp decode_call_status(%{"Committed" => _}), do: :committed
  defp decode_call_status(%{"Failed" => msg}), do: {:failed, msg}
  defp decode_call_status(%{"OutOfEnergy" => _}), do: :out_of_energy
  defp decode_call_status(_), do: nil

  defp decode_timestamp(nil), do: nil
  defp decode_timestamp(us) when is_integer(us), do: %Types.Timestamp{microseconds_since_epoch: us}

  defp decode_timestamp(%{"microseconds" => us}),
    do: %Types.Timestamp{microseconds_since_epoch: us}

  defp decode_timestamp(%{"__time_duration_micros__" => us}),
    do: %Types.Timestamp{microseconds_since_epoch: us}

  defp decode_message_id(nil), do: <<>>
  defp decode_message_id(list) when is_list(list), do: :binary.list_to_bin(list)
  defp decode_message_id(b64) when is_binary(b64), do: Base.decode64!(b64)

  defp get_query_id(%{"queryId" => %{"id" => id}}), do: id
  defp get_query_id(%{"queryId" => id}) when is_integer(id), do: id
  defp get_query_id(%{"query_id" => %{"id" => id}}), do: id
  defp get_query_id(%{"query_id" => id}) when is_integer(id), do: id
  defp get_query_id(_), do: nil
end