class Kafka::Producer

def assign_partitions!
  failed_messages = []
  topics_with_failures = Set.new
  @pending_message_queue.each do |message|
    partition = message.partition
    begin
      # If a message for a topic fails to receive a partition all subsequent
      # messages for the topic should be retried to preserve ordering
      if topics_with_failures.include?(message.topic)
        failed_messages << message
        next
      end
      if partition.nil?
        partition_count = @cluster.partitions_for(message.topic).count
        partition = Partitioner.partition_for_key(partition_count, message)
      end
      @buffer.write(
        value: message.value,
        key: message.key,
        topic: message.topic,
        partition: partition,
        create_time: message.create_time,
      )
    rescue Kafka::Error => e
      @instrumenter.instrument("topic_error.producer", {
        topic: message.topic,
        exception: [e.class.to_s, e.message],
      })
      topics_with_failures << message.topic
      failed_messages << message
    end
  end
  if failed_messages.any?
    failed_messages.group_by(&:topic).each do |topic, messages|
      @logger.error "Failed to assign partitions to #{messages.count} messages in #{topic}"
    end
    @cluster.mark_as_stale!
  end
  @pending_message_queue.replace(failed_messages)
end

def buffer_bytesize
  @pending_message_queue.bytesize + @buffer.bytesize
end

def buffer_messages
  messages = []
  @pending_message_queue.each do |message|
    messages << message
  end
  @buffer.each do |topic, partition, messages_for_partition|
    messages_for_partition.each do |message|
      messages << PendingMessage.new(
        message.value,
        message.key,
        topic,
        partition,
        nil,
        message.create_time
      )
    end
  end
  messages
end

def buffer_overflow(topic, message)
  @instrumenter.instrument("buffer_overflow.producer", {
    topic: topic,
  })
  raise BufferOverflow, message
end

def buffer_size
  @pending_message_queue.size + @buffer.size
end

def clear_buffer
  @buffer.clear
  @pending_message_queue.clear
end

def deliver_messages
  # There's no need to do anything if the buffer is empty.
  return if buffer_size == 0
  @instrumenter.instrument("deliver_messages.producer") do |notification|
    message_count = buffer_size
    notification[:message_count] = message_count
    notification[:attempts] = 0
    begin
      deliver_messages_with_retries(notification)
    ensure
      notification[:delivered_message_count] = message_count - buffer_size
    end
  end
end

def deliver_messages_with_retries(notification)
  attempt = 0
  @cluster.add_target_topics(@target_topics)
  operation = ProduceOperation.new(
    cluster: @cluster,
    buffer: @buffer,
    required_acks: @required_acks,
    ack_timeout: @ack_timeout,
    compressor: @compressor,
    logger: @logger,
    instrumenter: @instrumenter,
  )
  loop do
    attempt += 1
    notification[:attempts] = attempt
    begin
      @cluster.refresh_metadata_if_necessary!
    rescue ConnectionError => e
      raise DeliveryFailed.new(e, buffer_messages)
    end
    assign_partitions!
    operation.execute
    if @required_acks.zero?
      # No response is returned by the brokers, so we can't know which messages
      # have been successfully written. Our only option is to assume that they all
      # have.
      @buffer.clear
    end
    if buffer_size.zero?
      break
    elsif attempt <= @max_retries
      @logger.warn "Failed to send all messages; attempting retry #{attempt} of #{@max_retries} after #{@retry_backoff}s"
      sleep @retry_backoff
    else
      @logger.error "Failed to send all messages; keeping remaining messages in buffer"
      break
    end
  end
  unless @pending_message_queue.empty?
    # Mark the cluster as stale in order to force a cluster metadata refresh.
    @cluster.mark_as_stale!
    raise DeliveryFailed.new("Failed to assign partitions to #{@pending_message_queue.size} messages", buffer_messages)
  end
  unless @buffer.empty?
    partitions = @buffer.map {|topic, partition, _| "#{topic}/#{partition}" }.join(", ")
    raise DeliveryFailed.new("Failed to send messages to #{partitions}", buffer_messages)
  end
end

def initialize(cluster:, logger:, instrumenter:, compressor:, ack_timeout:, required_acks:, max_retries:, retry_backoff:, max_buffer_size:, max_buffer_bytesize:)
  @cluster = cluster
  @logger = logger
  @instrumenter = instrumenter
  @required_acks = required_acks == :all ? -1 : required_acks
  @ack_timeout = ack_timeout
  @max_retries = max_retries
  @retry_backoff = retry_backoff
  @max_buffer_size = max_buffer_size
  @max_buffer_bytesize = max_buffer_bytesize
  @compressor = compressor
  # The set of topics that are produced to.
  @target_topics = Set.new
  # A buffer organized by topic/partition.
  @buffer = MessageBuffer.new
  # Messages added by `#produce` but not yet assigned a partition.
  @pending_message_queue = PendingMessageQueue.new
end

def produce(value, key: nil, topic:, partition: nil, partition_key: nil, create_time: Time.now)
  message = PendingMessage.new(
    value && value.to_s,
    key && key.to_s,
    topic.to_s,
    partition && Integer(partition),
    partition_key && partition_key.to_s,
    create_time,
  )
  if buffer_size >= @max_buffer_size
    buffer_overflow topic,
      "Cannot produce to #{topic}, max buffer size (#{@max_buffer_size} messages) reached"
  end
  if buffer_bytesize + message.bytesize >= @max_buffer_bytesize
    buffer_overflow topic,
      "Cannot produce to #{topic}, max buffer bytesize (#{@max_buffer_bytesize} bytes) reached"
  end
  @target_topics.add(topic)
  @pending_message_queue.write(message)
  @instrumenter.instrument("produce_message.producer", {
    value: value,
    key: key,
    topic: topic,
    create_time: create_time,
    message_size: message.bytesize,
    buffer_size: buffer_size,
    max_buffer_size: @max_buffer_size,
  })
  nil
end

def shutdown
  @cluster.disconnect
end