module SupplyDrop
  module Plugin

    def rsync
      SupplyDrop::Util.thread_pool_size = puppet_parallel_rsync_pool_size
      servers = SupplyDrop::Util.optionally_async(find_servers_for_task(current_task), puppet_parallel_rsync)
      failed_servers = servers.map do |server|
        rsync_cmd = SupplyDrop::Rsync.command(
          puppet_source,
          SupplyDrop::Rsync.remote_address(server.user || fetch(:user, ENV['USER']), server.host, puppet_destination),
          :delete => true,
          :excludes => puppet_excludes,
          :ssh => ssh_options.merge(server.options[:ssh_options]||{}),
          :flags => '--copy-links'
        )
        logger.debug rsync_cmd
        server.host unless system rsync_cmd
      end.compact

      raise "rsync failed on #{failed_servers.join(',')}" if failed_servers.any?
    end

    def prepare
      #run "mkdir -p #{puppet_destination}"
      #run "#{sudo} chown -R $USER: #{puppet_destination}"
    end

    def noop
      puppet(:noop)
    end

    def apply
      puppet(:apply)
    end

    def lock
      if should_lock?
        cmd = <<-GETLOCK
if [ ! -f #{puppet_lock_file} ]; then
    touch #{puppet_lock_file};
else
    stat -c "#{red_text("Puppet in progress, #{puppet_lock_file} owned by %U since %x")}" #{puppet_lock_file} >&2;
    exit 1;
fi
        GETLOCK
        run cmd.gsub(/\s+/, ' ')
      end
    end

    def unlock
      run "rm -f #{puppet_lock_file}; true" if should_lock?
    end

    private

    def should_lock?
      puppet_lock_file && !ENV['NO_PUPPET_LOCK']
    end

    def puppet(command = :noop)
      #puppet_cmd = "cd #{puppet_destination} && #{sudo_cmd} #{puppet_command} --modulepath=#{puppet_lib} #{puppet_parameters}"
      puppet_cmd = "cd #{puppet_destination} && #{sudo_cmd} #{puppet_command} #{puppet_parameters}"
      flag = command == :noop ? '--noop' : ''

      writer = if puppet_stream_output
                 SupplyDrop::Writer::Streaming.new(logger)
               else
                 SupplyDrop::Writer::Batched.new(logger)
               end

      writer = SupplyDrop::Writer::File.new(writer, puppet_write_to_file) unless puppet_write_to_file.nil?

      begin
        exitcode = nil
        run "#{puppet_cmd} #{flag}; echo exitcode:$?" do |channel, stream, data|
          if data =~ /exitcode:(\d+)/
            exitcode = $1
            writer.collect_output(channel[:host], "Puppet #{command} complete (#{exitcode_description(exitcode)}).\n")
          else
            writer.collect_output(channel[:host], data)
          end
        end
      ensure
        writer.all_output_collected
      end
    end

    def exitcode_description(code)
      case code
        when "0" then "no changes"
        when "2" then "changes made"
        when "4" then "failed"
        when "6" then "changes and failures"
        else code
      end
    end

    def red_text(text)
      "\033[0;31m#{text}\033[0m"
    end

    def sudo_cmd
      if fetch(:use_sudo, true)
        sudo(:as => puppet_runner)
      else
        logger.info "NOTICE: puppet_runner configuration invalid when use_sudo is false, ignoring..." unless puppet_runner.nil?
        ''
      end
    end
  end
end