nebula_graph_ex

defmodule NebulaGraphEx.Thrift.GraphService do
  @moduledoc """
  Typed Thrift calls for the NebulaGraph `GraphService`.

  This module wraps `NebulaGraphEx.Thrift.Client` with the exact method
  signatures defined in `priv/thrift/graph.thrift`. All encoding and
  decoding is done manually against the Thrift binary protocol so that the
  raw socket stays under `DBConnection`'s control.

  ## Methods

  * `authenticate/3` — open a session with username + password
  * `execute/4` — run a nGQL statement (binary DataSet response)
  * `execute_with_parameter/5` — parameterised nGQL
  * `execute_json/4` — run a nGQL statement (JSON response)
  * `signout/2` — close a session (one-way, no reply)
  * `verify_client_version/2` — version handshake
  """

  alias NebulaGraphEx.Thrift.Client
  alias NebulaGraphEx.Thrift.Types
  alias NebulaGraphEx.Transport

  # Thrift type constants — must be defined before any function that
  # uses them in pattern matching.
  @ti32 Client.type_i32()
  @ti64 Client.type_i64()
  @ti16 Client.type_i16()
  @tbyte Client.type_byte()
  @tbool Client.type_bool()
  @tdouble Client.type_double()
  @tbinary Client.type_binary()
  @tstruct Client.type_struct()
  @tlist Client.type_list()
  @tset Client.type_set()
  @tmap Client.type_map()

  @nebula_version "3.0.0"

  # ─── authenticate ──────────────────────────────────────────────────────────

  @doc """
  Authenticates against the NebulaGraph server.

  Returns `{:ok, auth_response}` on success, where `auth_response` is a map:

  * `:error_code` — raw integer error code (0 = success)
  * `:session_id` — integer session ID to use in subsequent calls
  * `:time_zone_offset_seconds` — server timezone offset or `nil`
  * `:time_zone_name` — server timezone binary or `nil`
  * `:error_msg` — error message binary or `nil`
  """
  @spec authenticate(Transport.socket(), binary(), binary(), keyword()) ::
          {:ok, map()} | {:error, term()}
  def authenticate(socket, username, password, opts \\ []) do
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

    args = [
      Client.field(@tbinary, 1),
      Client.binary_val(username),
      Client.field(@tbinary, 2),
      Client.binary_val(password),
      Client.stop()
    ]

    with {:ok, body} <- Client.call(socket, "authenticate", args, 1, recv_timeout) do
      decode_auth_response(body)
    end
  end

  # ─── signout ───────────────────────────────────────────────────────────────

  @doc """
  Signs out and closes the session. One-way — the server sends no reply.
  """
  @spec signout(Transport.socket(), integer()) :: :ok | {:error, term()}
  def signout(socket, session_id) do
    args = [
      Client.field(@ti64, 1),
      Client.i64(session_id),
      Client.stop()
    ]

    Client.cast(socket, "signout", args)
  end

  # ─── execute ───────────────────────────────────────────────────────────────

  @doc """
  Executes a nGQL statement and returns an `ExecutionResponse` map with keys:

  * `:error_code` — integer
  * `:latency_in_us` — server-side latency in microseconds
  * `:data` — `nil` or a map with `:column_names` and `:rows`
  * `:space_name` — binary or `nil`
  * `:error_msg` — binary or `nil`
  * `:comment` — binary or `nil`
  """
  @spec execute(Transport.socket(), integer(), binary(), keyword()) ::
          {:ok, map()} | {:error, term()}
  def execute(socket, session_id, stmt, opts \\ []) do
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

    args = [
      Client.field(@ti64, 1),
      Client.i64(session_id),
      Client.field(@tbinary, 2),
      Client.binary_val(stmt),
      Client.stop()
    ]

    with {:ok, body} <- Client.call(socket, "execute", args, 2, recv_timeout) do
      decode_execution_response(body)
    end
  end

  # ─── execute_with_parameter ────────────────────────────────────────────────

  @doc """
  Executes a parameterised nGQL statement.

  `params` is a map with binary keys and Elixir values that are encodable
  by `NebulaGraphEx.Thrift.Types.encode_value/1`.

  ## Example

      GraphService.execute_with_parameter(
        socket,
        session_id,
        "MATCH (v:Player{name: $name}) RETURN v",
        %{"name" => "Tim Duncan"}
      )
  """
  @spec execute_with_parameter(
          Transport.socket(),
          integer(),
          binary(),
          map(),
          keyword()
        ) :: {:ok, map()} | {:error, term()}
  def execute_with_parameter(socket, session_id, stmt, params, opts \\ []) do
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

    encoded_params =
      Enum.map(params, fn {k, v} ->
        [Client.binary_val(k), Types.encode_value(v)]
      end)

    args = [
      Client.field(@ti64, 1),
      Client.i64(session_id),
      Client.field(@tbinary, 2),
      Client.binary_val(stmt),
      Client.field(@tmap, 3),
      Client.map(@tbinary, @tstruct, encoded_params),
      Client.stop()
    ]

    with {:ok, body} <- Client.call(socket, "executeWithParameter", args, 3, recv_timeout) do
      decode_execution_response(body)
    end
  end

  # ─── execute_json ──────────────────────────────────────────────────────────

  @doc """
  Executes a nGQL statement and returns a raw JSON binary.

  Useful for debugging. The caller is responsible for JSON decoding.
  """
  @spec execute_json(Transport.socket(), integer(), binary(), keyword()) ::
          {:ok, binary()} | {:error, term()}
  def execute_json(socket, session_id, stmt, opts \\ []) do
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

    args = [
      Client.field(@ti64, 1),
      Client.i64(session_id),
      Client.field(@tbinary, 2),
      Client.binary_val(stmt),
      Client.stop()
    ]

    with {:ok, body} <- Client.call(socket, "executeJson", args, 4, recv_timeout) do
      decode_binary_return(body)
    end
  end

  # ─── verify_client_version ─────────────────────────────────────────────────

  @doc """
  Performs the client version handshake. Call immediately after opening
  the socket, before `authenticate/3`.

  Returns `{:ok, :ok}` on success or `{:error, message}`.
  """
  @spec verify_client_version(Transport.socket(), keyword()) ::
          {:ok, :ok} | {:error, term()}
  def verify_client_version(socket, opts \\ []) do
    recv_timeout = Keyword.get(opts, :recv_timeout, 15_000)

    req_fields = [
      Client.field(@tbinary, 1),
      Client.binary_val(@nebula_version),
      Client.stop()
    ]

    args = [
      Client.field(@tstruct, 1),
      req_fields,
      Client.stop()
    ]

    with {:ok, body} <- Client.call(socket, "verifyClientVersion", args, 5, recv_timeout) do
      decode_verify_response(body)
    end
  end

  # ─── Response decoders ─────────────────────────────────────────────────────

  defp decode_auth_response(body) do
    inner =
      case Client.next_field(body) do
        {:ok, 0, @tstruct, rest} -> rest
        _ -> body
      end

    decode_struct(
      inner,
      %{
        error_code: nil,
        error_msg: nil,
        session_id: nil,
        time_zone_offset_seconds: nil,
        time_zone_name: nil
      },
      fn field_id, type, rest, acc ->
        case {field_id, type} do
          {1, @ti32} ->
            {:ok, v, rest2} = Client.decode_i32(rest)
            {Map.put(acc, :error_code, v), rest2}

          {2, @tbinary} ->
            {:ok, v, rest2} = Client.decode_string(rest)
            {Map.put(acc, :error_msg, v), rest2}

          {3, @ti64} ->
            {:ok, v, rest2} = Client.decode_i64(rest)
            {Map.put(acc, :session_id, v), rest2}

          {4, @ti32} ->
            {:ok, v, rest2} = Client.decode_i32(rest)
            {Map.put(acc, :time_zone_offset_seconds, v), rest2}

          {5, @tbinary} ->
            {:ok, v, rest2} = Client.decode_string(rest)
            {Map.put(acc, :time_zone_name, v), rest2}

          _ ->
            {acc, skip_field(type, rest)}
        end
      end
    )
  end

  defp decode_execution_response(body) do
    inner =
      case Client.next_field(body) do
        {:ok, 0, @tstruct, rest} -> rest
        _ -> body
      end

    decode_struct(
      inner,
      %{
        error_code: nil,
        latency_in_us: nil,
        data: nil,
        space_name: nil,
        error_msg: nil,
        comment: nil
      },
      fn field_id, type, rest, acc ->
        case {field_id, type} do
          {1, @ti32} ->
            {:ok, v, rest2} = Client.decode_i32(rest)
            {Map.put(acc, :error_code, v), rest2}

          {2, @ti64} ->
            {:ok, v, rest2} = Client.decode_i64(rest)
            {Map.put(acc, :latency_in_us, v), rest2}

          {3, @tstruct} ->
            {dataset, rest2} = Types.decode_dataset(rest)
            {Map.put(acc, :data, dataset), rest2}

          {4, @tbinary} ->
            {:ok, v, rest2} = Client.decode_string(rest)
            {Map.put(acc, :space_name, v), rest2}

          {5, @tbinary} ->
            {:ok, v, rest2} = Client.decode_string(rest)
            {Map.put(acc, :error_msg, v), rest2}

          {7, @tbinary} ->
            {:ok, v, rest2} = Client.decode_string(rest)
            {Map.put(acc, :comment, v), rest2}

          _ ->
            {acc, skip_field(type, rest)}
        end
      end
    )
  end

  defp decode_verify_response(body) do
    # Thrift reply body: field_id=0 (struct) wraps the actual return value.
    inner =
      case Client.next_field(body) do
        {:ok, 0, @tstruct, rest} -> rest
        _ -> body
      end

    with {:ok, resp} <-
           decode_struct(
             inner,
             %{error_code: nil, error_msg: nil},
             fn field_id, type, rest, acc ->
               case {field_id, type} do
                 {1, @ti32} ->
                   {:ok, v, rest2} = Client.decode_i32(rest)
                   {Map.put(acc, :error_code, v), rest2}

                 {2, @tbinary} ->
                   {:ok, v, rest2} = Client.decode_string(rest)
                   {Map.put(acc, :error_msg, v), rest2}

                 _ ->
                   {acc, skip_field(type, rest)}
               end
             end
           ) do
      case resp do
        %{error_code: 0} -> {:ok, :ok}
        %{error_msg: msg} when is_binary(msg) -> {:error, msg}
        _ -> {:error, :verify_failed}
      end
    end
  end

  defp decode_binary_return(body) do
    case Client.next_field(body) do
      {:ok, _id, @tbinary, rest} -> Client.decode_string(rest)
      {:stop, _} -> {:ok, nil}
      _ -> {:error, :unexpected_reply_format}
    end
  end

  # ─── Generic struct decoder ────────────────────────────────────────────────

  defp decode_struct(data, initial_acc, handler) do
    do_decode_struct(data, initial_acc, handler)
  end

  defp do_decode_struct(data, acc, handler) do
    case Client.next_field(data) do
      {:stop, _rest} ->
        {:ok, acc}

      {:ok, field_id, type, rest} ->
        {new_acc, rest2} = handler.(field_id, type, rest, acc)
        do_decode_struct(rest2, new_acc, handler)

      other ->
        {:error, {:decode_error, other}}
    end
  end

  # ─── Field skip (tolerates unknown/future fields) ──────────────────────────

  defp skip_field(@ti32, <<_::32, rest::binary>>), do: rest
  defp skip_field(@ti64, <<_::64, rest::binary>>), do: rest
  defp skip_field(@ti16, <<_::16, rest::binary>>), do: rest
  defp skip_field(@tbyte, <<_::8, rest::binary>>), do: rest
  defp skip_field(@tbool, <<_::8, rest::binary>>), do: rest
  defp skip_field(@tdouble, <<_::64, rest::binary>>), do: rest

  defp skip_field(@tbinary, <<len::32-big, _::binary-size(len), rest::binary>>), do: rest

  defp skip_field(@tstruct, data), do: skip_struct_fields(data)

  defp skip_field(@tlist, <<etype::8, count::32-big, rest::binary>>),
    do: skip_n(etype, count, rest)

  defp skip_field(@tset, <<etype::8, count::32-big, rest::binary>>),
    do: skip_n(etype, count, rest)

  defp skip_field(@tmap, <<ktype::8, vtype::8, count::32-big, rest::binary>>) do
    Enum.reduce(1..max(count, 1), rest, fn _, r ->
      r |> skip_field(ktype) |> skip_field(vtype)
    end)
  end

  defp skip_field(_, rest), do: rest

  defp skip_struct_fields(data) do
    case Client.next_field(data) do
      {:stop, rest} -> rest
      {:ok, _id, type, rest} -> skip_struct_fields(skip_field(type, rest))
    end
  end

  defp skip_n(_type, 0, rest), do: rest

  defp skip_n(type, n, rest) do
    rest |> skip_field(type) |> skip_n(type, n - 1)
  end
end