mb_recipe :aptitude do
  during :provision, %w[check upgrade install]
end

namespace :mb do
  namespace :aptitude do
    desc "Verify server is Ubuntu 16.04 or 18.04"
    task :check do
      privileged_on roles(:all) do
        version = capture(:sudo, "lsb_release -a")[/^Release:\s+(\S+)$/, 1]
        next if %w[16.04 18.04].include?(version)

        raise "Ubuntu version #{version || "unknown"} is not supported by "\
              "capistrano-mb. Only Ubuntu 16.04 and 18.04 are supported. "\
              "Downgrade capistrano-mb if you need to use an older version of "\
              "Ubuntu."
      end
    end

    desc "Run `apt update` and then run `apt upgrade`"
    task :upgrade do
      privileged_on roles(:all) do
        _update
        _upgrade
      end
    end

    desc "Run `apt install` for packages required by the roles of "\
         "each server."
    task :install do
      privileged_on roles(:all) do |host|
        packages_to_install = []
        repos_to_add = []

        _each_package(host) do |pkg, repo|
          unless _already_installed?(pkg)
            repos_to_add << repo unless repo.nil?
            packages_to_install << pkg
          end
        end

        repos_to_add.uniq.each { |repo| _add_repository(repo) }
        _update
        packages_to_install.uniq.each { |pkg| _install(pkg) }
      end
    end

    def _already_installed?(pkg)
      test(:sudo,
           "dpkg", "-s", pkg,
           "2>/dev/null", "|", :grep, "-q 'ok installed'")
    end

    def _add_repository(repo)
      unless _already_installed?("software-properties-common")
        _install("software-properties-common")
      end
      execute :sudo, "apt-add-repository", "-y '#{repo}'"
    end

    def _install(pkg)
      execute :sudo, "DEBIAN_FRONTEND=noninteractive apt-get -y install", pkg
    end

    def _update
      execute :sudo, "DEBIAN_FRONTEND=noninteractive apt-get -y update"
    end

    def _upgrade
      execute :sudo,
              "DEBIAN_FRONTEND=noninteractive apt-get -y "\
              '-o DPkg::options::="--force-confdef" '\
              '-o DPkg::options::="--force-confold" '\
              "upgrade"
    end

    def _each_package(host)
      return to_enum(:_each_package, host) unless block_given?
      hostname = host.hostname

      fetch(:mb_aptitude_packages).each do |package_spec, *role_list|
        next unless roles(*role_list.flatten).map(&:hostname).include?(hostname)

        pkg, repo = package_spec.split("@")
        yield(pkg, repo)
      end
    end
  end
end
