module CustomAirbrusshFormatter
  def initialize(io, config)
    @line_open = false
    @command_had_output = {}.compare_by_identity
    @command_stderr = {}.compare_by_identity
    super
  end

  def log_command_start(command)
    return if debug?(command)

    @command_had_output[command] = false
    @command_stderr[command] = +''
    print_task_if_changed
  end

  def log_command_data(command, stream_type, string)
    return if debug?(command)

    data = string.to_s
    return if data.strip.empty?

    @command_stderr[command] ||= +''
    @command_stderr[command] << data if stream_type == :stderr
    print_progress_dot
    @command_had_output[command] = true
  end

  def log_command_exit(command)
    return if debug?(command)

    if command.success?
      print_progress_dot unless @command_had_output[command]
    else
      ensure_line_break_if_needed
      print_buffered_stderr(command)
      super
    end
  ensure
    @command_had_output.delete(command)
    @command_stderr.delete(command)
  end

  def write_log_message(log_message)
    return if debug?(log_message)
    return if log_message.verbosity < SSHKit::Logger::WARN

    super
  end

  private

  def print_task_if_changed
    return if current_task_name.nil?
    return if current_task_name == last_printed_task

    ensure_line_break_if_needed
    self.last_printed_task = current_task_name
    @console << "#{clock} #{blue(current_task_name)} "
    @line_open = true
  end

  def print_progress_dot
    @console << green('.')
    @line_open = true
  end

  def ensure_line_break_if_needed
    return unless @line_open

    @console << "\n"
    @line_open = false
  end

  def print_buffered_stderr(command)
    buffered_stderr = @command_stderr[command].to_s
    return if buffered_stderr.strip.empty?

    buffered_stderr.each_line do |line|
      stripped_line = line.rstrip
      next if stripped_line.empty?

      @console << red(stripped_line) << "\n"
    end
  end
end

Airbrussh::ConsoleFormatter.prepend CustomAirbrusshFormatter