class Selenium::WebDriver::BiDi::LogHandler

def add_message_handler(type)
  subscribe_log_entry unless @log_entry_subscribed
  @bidi.add_callback('log.entryAdded') do |params|
    if params['type'] == type
      log_entry_klass = type == 'console' ? ConsoleLogEntry : JavaScriptLogEntry
      yield(log_entry_klass.new(**params))
    end
  end
end

def initialize(bidi)
  @bidi = bidi
  @log_entry_subscribed = false
end

def remove_message_handler(id)
  @bidi.remove_callback('log.entryAdded', id)
  unsubscribe_log_entry if @log_entry_subscribed && @bidi.callbacks['log.entryAdded'].empty?
end

def subscribe_log_entry
  @bidi.session.subscribe('log.entryAdded')
  @log_entry_subscribed = true
end

def unsubscribe_log_entry
  @bidi.session.unsubscribe('log.entryAdded')
  @log_entry_subscribed = false
end