ex_kits

defmodule ExKits.Utils.BufferV2 do
  @moduledoc """
  A high-performance, fault-tolerant buffer queue implementation using GenServer.

  This module provides a way to create a buffer queue with a specified capacity
  and allows items to be put into and taken out of the queue. Unlike the original
  Buffer implementation, BufferV2 uses GenServer for proper blocking semantics
  and better concurrent performance.

  ## Key Features

  - **Blocking and non-blocking operations**: `take/1,2` blocks until data is available,
    `try_take/1` returns immediately
  - **Timeout support**: All blocking operations support configurable timeouts
  - **Batch operations**: Support for putting/taking multiple items at once
  - **Memory efficient**: Uses `:queue` for O(1) operations and stable memory usage
  - **OTP compliant**: Built with GenServer for proper supervision and fault tolerance

  ## Examples

      # Start a named buffer
      {:ok, pid} = ExKits.Utils.BufferV2.start_link(name: :my_buffer, capacity: 1000)

      # Put items (non-blocking)
      :ok = ExKits.Utils.BufferV2.put(:my_buffer, :item1)
      :ok = ExKits.Utils.BufferV2.put_batch(:my_buffer, [:item2, :item3])

      # Take items (blocking)
      {:ok, :item1} = ExKits.Utils.BufferV2.take(:my_buffer)
      {:ok, :item2} = ExKits.Utils.BufferV2.take(:my_buffer, 5000)  # 5 second timeout

      # Take items (non-blocking)
      {:ok, :item3} = ExKits.Utils.BufferV2.try_take(:my_buffer)
      {:error, :empty} = ExKits.Utils.BufferV2.try_take(:my_buffer)

      # Batch operations
      {:ok, items} = ExKits.Utils.BufferV2.take_batch(:my_buffer, 3)
      {:ok, items} = ExKits.Utils.BufferV2.try_take_batch(:my_buffer, 5)

  """

  use GenServer

  require Logger

  @type buffer_name :: atom() | pid()
  @type timeout_ms :: non_neg_integer() | :infinity

  defmodule State do
    @moduledoc false

    @type t :: %__MODULE__{
            name: atom() | nil,
            capacity: non_neg_integer(),
            queue: :queue.queue(),
            size: non_neg_integer(),
            waiting_consumers: :queue.queue()
          }

    defstruct name: nil,
              capacity: 0,
              queue: :queue.new(),
              size: 0,
              waiting_consumers: :queue.new()
  end

  defmodule Consumer do
    @moduledoc false

    @type t :: %__MODULE__{
            from: GenServer.from(),
            count: pos_integer(),
            timeout_ref: reference()
          }

    defstruct from: nil,
              count: 1,
              timeout_ref: nil
  end

  #
  # Public API
  #

  @doc """
  Starts a new BufferV2 process.

  ## Options

  * `:name` - (optional) The name to register the buffer process
  * `:capacity` - The maximum number of items the buffer can hold

  ## Examples

      # Named process
      {:ok, pid} = BufferV2.start_link(name: :my_buffer, capacity: 1000)

      # Anonymous process
      {:ok, pid} = BufferV2.start_link(capacity: 500)

  """
  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts) do
    {name, opts} = Keyword.pop(opts, :name)
    {capacity, opts} = Keyword.pop!(opts, :capacity)

    if capacity <= 0 do
      raise ArgumentError, "capacity must be a positive integer, got: #{inspect(capacity)}"
    end

    state = %State{
      name: name,
      capacity: capacity
    }

    if name do
      GenServer.start_link(__MODULE__, state, [name: name] ++ opts)
    else
      GenServer.start_link(__MODULE__, state, opts)
    end
  end

  @doc """
  Stops the buffer process gracefully.

  ## Examples

      :ok = BufferV2.stop(:my_buffer)
      :ok = BufferV2.stop(pid)

  """
  @spec stop(buffer_name()) :: :ok
  def stop(name_or_pid) do
    GenServer.stop(name_or_pid)
  end

  @doc """
  Puts a single item into the buffer.

  Returns `:ok` if successful, or `{:error, :full}` if the buffer is at capacity.
  This operation is non-blocking.

  ## Examples

      :ok = BufferV2.put(:my_buffer, :item)
      {:error, :full} = BufferV2.put(:my_buffer, :item)

  """
  @spec put(buffer_name(), any()) :: :ok | {:error, :full}
  def put(name_or_pid, item) do
    GenServer.call(name_or_pid, {:put, item})
  end

  @doc """
  Puts multiple items into the buffer.

  Returns `:ok` if all items were added successfully, or `{:error, :full}` if
  there isn't enough space for all items. This operation is atomic - either
  all items are added or none are.

  ## Examples

      :ok = BufferV2.put_batch(:my_buffer, [:item1, :item2, :item3])
      {:error, :full} = BufferV2.put_batch(:my_buffer, too_many_items)

  """
  @spec put_batch(buffer_name(), [any()]) :: :ok | {:error, :full}
  def put_batch(name_or_pid, items) when is_list(items) do
    GenServer.call(name_or_pid, {:put_batch, items})
  end

  @doc """
  Takes a single item from the buffer, blocking until one is available.

  Returns `{:ok, item}` when an item is available, or `{:error, :timeout}`
  if the timeout is reached.

  ## Examples

      {:ok, item} = BufferV2.take(:my_buffer)              # Block indefinitely
      {:ok, item} = BufferV2.take(:my_buffer, 5000)        # 5 second timeout
      {:error, :timeout} = BufferV2.take(:my_buffer, 1000) # Timeout

  """
  @spec take(buffer_name(), timeout_ms()) :: {:ok, any()} | {:error, :timeout}
  def take(name_or_pid, timeout \\ :infinity) do
    GenServer.call(name_or_pid, {:take, 1, timeout}, get_call_timeout(timeout))
  end

  @doc """
  Takes multiple items from the buffer, blocking until the requested count is available.

  Returns `{:ok, items}` when enough items are available, or `{:error, :timeout}`
  if the timeout is reached.

  ## Examples

      {:ok, [item1, item2]} = BufferV2.take_batch(:my_buffer, 2)
      {:ok, items} = BufferV2.take_batch(:my_buffer, 5, 3000)  # 3 second timeout
      {:error, :timeout} = BufferV2.take_batch(:my_buffer, 10, 1000)

  """
  @spec take_batch(buffer_name(), pos_integer(), timeout_ms()) :: {:ok, [any()]} | {:error, :timeout}
  def take_batch(name_or_pid, count, timeout \\ :infinity) when count > 0 do
    GenServer.call(name_or_pid, {:take, count, timeout}, get_call_timeout(timeout))
  end

  @doc """
  Tries to take a single item from the buffer without blocking.

  Returns `{:ok, item}` if an item is immediately available, or `{:error, :empty}`
  if the buffer is empty.

  ## Examples

      {:ok, item} = BufferV2.try_take(:my_buffer)
      {:error, :empty} = BufferV2.try_take(:my_buffer)

  """
  @spec try_take(buffer_name()) :: {:ok, any()} | {:error, :empty}
  def try_take(name_or_pid) do
    GenServer.call(name_or_pid, {:try_take, 1})
  end

  @doc """
  Tries to take multiple items from the buffer without blocking.

  Returns `{:ok, items}` with as many items as are immediately available (up to count),
  or `{:ok, []}` if the buffer is empty.

  ## Examples

      {:ok, [item1, item2]} = BufferV2.try_take_batch(:my_buffer, 5)  # Got 2 out of 5
      {:ok, []} = BufferV2.try_take_batch(:my_buffer, 3)             # Buffer empty

  """
  @spec try_take_batch(buffer_name(), pos_integer()) :: {:ok, [any()]}
  def try_take_batch(name_or_pid, count) when count > 0 do
    GenServer.call(name_or_pid, {:try_take, count})
  end

  @doc """
  Returns information about the buffer state.

  ## Examples

      %{size: 5, capacity: 1000, waiting_consumers: 2} = BufferV2.info(:my_buffer)

  """
  @spec info(buffer_name()) :: %{
          size: non_neg_integer(),
          capacity: non_neg_integer(),
          waiting_consumers: non_neg_integer()
        }
  def info(name_or_pid) do
    GenServer.call(name_or_pid, :info)
  end

  #
  # GenServer Callbacks
  #

  @impl GenServer
  def init(%State{} = state) do
    Logger.debug("BufferV2 started with capacity #{state.capacity}")
    {:ok, state}
  end

  @impl GenServer
  def handle_call({:put, item}, _from, %State{} = state) do
    handle_put([item], state)
  end

  def handle_call({:put_batch, items}, _from, %State{} = state) do
    handle_put(items, state)
  end

  def handle_call({:take, count, timeout}, from, %State{} = state) do
    handle_take(count, timeout, from, state)
  end

  def handle_call({:try_take, count}, _from, %State{} = state) do
    handle_try_take(count, state)
  end

  def handle_call(:info, _from, %State{} = state) do
    info = %{
      size: state.size,
      capacity: state.capacity,
      waiting_consumers: :queue.len(state.waiting_consumers)
    }

    {:reply, info, state}
  end

  @impl GenServer
  def handle_info({:timeout, ref, :consumer_timeout}, %State{} = state) do
    handle_consumer_timeout(ref, state)
  end

  def handle_info(msg, state) do
    Logger.warning("BufferV2 received unexpected message: #{inspect(msg)}")
    {:noreply, state}
  end

  @impl GenServer
  def terminate(reason, %State{} = state) do
    Logger.debug("BufferV2 terminating: #{inspect(reason)}")

    # Reply to all waiting consumers with error
    state.waiting_consumers
    |> :queue.to_list()
    |> Enum.each(fn %Consumer{from: from, timeout_ref: timeout_ref} ->
      if timeout_ref, do: Process.cancel_timer(timeout_ref)
      GenServer.reply(from, {:error, :shutdown})
    end)

    :ok
  end

  #
  # Internal Functions
  #

  defp handle_put(items, %State{size: size, capacity: capacity} = state) when is_list(items) do
    items_count = length(items)

    if size + items_count > capacity do
      {:reply, {:error, :full}, state}
    else
      # Add items to queue
      new_queue = Enum.reduce(items, state.queue, fn item, acc -> :queue.in(item, acc) end)
      new_state = %State{state | queue: new_queue, size: size + items_count}

      # Try to satisfy waiting consumers
      final_state = satisfy_waiting_consumers(new_state)

      {:reply, :ok, final_state}
    end
  end

  defp handle_take(count, timeout, from, %State{size: size} = state) do
    if size >= count do
      # Enough items available, return immediately
      {items, new_state} = take_items(count, state)
      result = if count == 1, do: {:ok, List.first(items)}, else: {:ok, items}
      {:reply, result, new_state}
    else
      # Not enough items, add to waiting queue
      timeout_ref =
        if timeout != :infinity do
          ref = make_ref()
          Process.send_after(self(), {:timeout, ref, :consumer_timeout}, timeout)
          ref
        end

      consumer = %Consumer{
        from: from,
        count: count,
        timeout_ref: timeout_ref
      }

      new_waiting = :queue.in(consumer, state.waiting_consumers)
      new_state = %State{state | waiting_consumers: new_waiting}

      {:noreply, new_state}
    end
  end

  defp handle_try_take(count, %State{} = state) do
    available = min(count, state.size)

    if available == 0 do
      if count == 1 do
        {:reply, {:error, :empty}, state}
      else
        {:reply, {:ok, []}, state}
      end
    else
      {items, new_state} = take_items(available, state)
      result = if count == 1, do: {:ok, List.first(items)}, else: {:ok, items}
      {:reply, result, new_state}
    end
  end

  defp handle_consumer_timeout(timeout_ref, %State{} = state) do
    {new_waiting, found_consumer} = remove_consumer_by_timeout_ref(state.waiting_consumers, timeout_ref)

    case found_consumer do
      %Consumer{from: from} ->
        GenServer.reply(from, {:error, :timeout})
        {:noreply, %State{state | waiting_consumers: new_waiting}}

      nil ->
        # Consumer already processed, ignore timeout
        {:noreply, state}
    end
  end

  defp satisfy_waiting_consumers(%State{waiting_consumers: waiting} = state) do
    case :queue.out(waiting) do
      {{:value, %Consumer{count: count} = consumer}, remaining_waiting} ->
        if state.size >= count do
          # Satisfy this consumer
          {items, new_state} = take_items(count, state)

          if consumer.timeout_ref do
            Process.cancel_timer(consumer.timeout_ref)
          end

          result = if count == 1, do: {:ok, List.first(items)}, else: {:ok, items}
          GenServer.reply(consumer.from, result)

          # Update state and try to satisfy more consumers
          updated_state = %State{new_state | waiting_consumers: remaining_waiting}
          satisfy_waiting_consumers(updated_state)
        else
          # Can't satisfy this consumer yet, stop processing
          state
        end

      {:empty, _} ->
        # No waiting consumers
        state
    end
  end

  defp take_items(count, %State{queue: queue, size: size} = state) do
    {items, new_queue} = take_from_queue(queue, count)
    new_state = %State{state | queue: new_queue, size: size - count}
    {items, new_state}
  end

  # Optimized queue extraction using :queue.split for O(1) performance
  #
  # This implementation provides significant performance improvements over
  # recursive item-by-item extraction:
  # - :queue.split is O(1) operation vs O(n) recursive extraction
  # - Eliminates function call overhead from recursion
  # - Reduces memory allocation by avoiding accumulator reversal
  # - Particularly beneficial for large batch operations
  defp take_from_queue(queue, count) do
    queue_length = :queue.len(queue)
    actual_count = min(count, queue_length)

    if actual_count == 0 do
      {[], queue}
    else
      {front_queue, back_queue} = :queue.split(actual_count, queue)
      items = :queue.to_list(front_queue)
      {items, back_queue}
    end
  end

  defp remove_consumer_by_timeout_ref(waiting_queue, target_ref) do
    remove_consumer_by_timeout_ref(:queue.out(waiting_queue), target_ref, :queue.new())
  end

  defp remove_consumer_by_timeout_ref({:empty, _}, _target_ref, acc) do
    {acc, nil}
  end

  defp remove_consumer_by_timeout_ref({{:value, %Consumer{timeout_ref: ref} = consumer}, remaining}, target_ref, acc) do
    if ref == target_ref do
      # Found the consumer, return the queue without this consumer
      new_queue = :queue.join(acc, remaining)
      {new_queue, consumer}
    else
      # Not the target consumer, keep looking
      new_acc = :queue.in(consumer, acc)
      remove_consumer_by_timeout_ref(:queue.out(remaining), target_ref, new_acc)
    end
  end

  defp get_call_timeout(:infinity), do: :infinity
  defp get_call_timeout(timeout) when is_integer(timeout) and timeout > 0, do: timeout + 1000
end