defmodule Pulsarix do
  @moduledoc """
  Pulsarix is responsible for sending messages to Kafka topics.
  It supports various message production methods, including synchronous,
  asynchronous, and callback-based approaches.
  """

  @telemetry_produce_event [:pulsarix, :produce, :msg]

  @typedoc """
  Unix time in milliseconds.
  """
  @type msg_ts() :: non_neg_integer()

  @type key() :: binary()
  @type value() :: binary()
  @type headers() :: [{binary(), binary()}]

  @type msg_input() :: %{
          key: key(),
          ts: msg_ts(),
          value: value(),
          headers: headers()
        }
  @type batch_input() :: [msg_input()]

  @type data() ::
          value()
          | msg_input()
          | batch_input()
          | [{key(), value()}]
          | {msg_ts(), value()}
          | [{msg_ts(), key(), value()}]

  @type topic() :: binary()
  @type client() :: atom()

  @type offset() :: non_neg_integer()
  @type partition() :: non_neg_integer()
  @type produce_ack_cb() :: (partition(), offset() -> any())

  ## APIs

  @doc """
  Produces a message to Kafka asynchronously. The function returns immediately
  after buffering the message in the producer process.

  ## Example

      iex> Pulsarix.Producer.produce(:client, "topic", "key", "value")
      :ok

  """
  @spec produce(client(), topic(), key(), data()) :: :ok | {:error, term()}
  def produce(client, topic, key, data) do
    do_produce(:produce, client, topic, key, data)
  end

  @doc """
  Sync version of `produce/4`. This function will not return until a response
  is received from kafka, however if producer is started with `required_acks`
  set to `0`, this function will return once the messages are buffered in
  the producer process.

  ## Example

      iex> Pulsarix.Producer.produce_sync(:client, "topic", "key", "value")
      :ok
  """
  @spec produce_sync(client(), topic(), key(), data()) :: :ok | {:error, term()}
  def produce_sync(client, topic, key, data) do
    do_produce(:produce_sync, client, topic, key, data)
  end

  @doc """
  Same as `produce/4`, only the ack is not delivered as a message, instead,
  the callback is evaluated by producer worker when ack is received from kafka.

  ## Example

      iex> Pulsarix.Producer.produce_async(:client, "topic", "key", "value", fn _, _ -> :ok end)
      :ok
  """
  @spec produce_async(client(), topic(), key(), data(), produce_ack_cb()) ::
          :ok | {:error, term()}
  def produce_async(client, topic, key, data, callback) when is_function(callback, 2) do
    do_produce(:produce_cb, client, topic, key, data, [callback])
  end

  @doc """
  Produces a message to Kafka without waiting for any acknowledgment.

  ## Example

      iex> Pulsarix.Producer.produce_no_ack(:client, "topic", "key", "value")
      :ok
  """
  @spec produce_no_ack(client(), topic(), key(), data()) :: :ok | {:error, term()}
  def produce_no_ack(client, topic, key, data) do
    do_produce(:produce_no_ack, client, topic, key, data)
  end

  ## Internal functions

  defp get_partition(client, topic, key) do
    {:ok, partitions} = :brod.get_partitions_count_safe(client, topic)

    :erlang.phash2(key, partitions)
  end

  defp do_produce(type, client, topic, key, data, callback \\ []) do
    partition = get_partition(client, topic, key)

    :telemetry.span(@telemetry_produce_event, %{}, fn ->
      :brod
      |> apply(type, [client, topic, partition, key, data | callback])
      |> span_result(%{key: key, data: data, topic: topic, client: client})
    end)
  end

  defp span_result(result, metadadata) do
    case result do
      {:error, error} ->
        {{:error, error}, Map.merge(metadadata, %{status: :error, reason: error})}

      {:ok, result} ->
        {:ok, Map.merge(metadadata, %{status: :success, result: result})}

      :ok ->
        {:ok, Map.merge(metadadata, %{status: :success, result: :ok})}
    end
  end
end