require 'cgi'
require 'beaker/dsl/install_utils'
require 'beaker/dsl/helpers'
require 'beaker-puppet'
require 'open3'
require 'pp'
require 'set'
require 'test/unit/assertions'
require 'json'
require 'inifile'
require 'mkmf' # provides find_executable()

module PuppetDBExtensions
  include Test::Unit::Assertions

  GitReposDir = Beaker::DSL::InstallUtils::FOSSUtils::SourcePath

  # We include the Puppet path here so we can use the Puppet ecosystem when
  # using rake, which needs the facter gem to work - AIO ruby contains this
  # by default.
  LeinCommandPrefix = "cd #{GitReposDir}/puppetdb; LEIN_ROOT=true LEIN_SNAPSHOTS_IN_RELEASE=true PATH=/opt/puppetlabs/puppet/bin:$PATH"

  def self.initialize_test_config(options, os_families)

    base_dir = File.join(File.dirname(__FILE__), '..')

    install_type = get_option_value(options[:puppetdb_install_type],
      [:git, :package, :pe], "install type", "PUPPETDB_INSTALL_TYPE", :package)

    install_mode =
        get_option_value(options[:install_type],
                         [:install, :upgrade_latest, :upgrade_oldest], "install type",
                         "INSTALL_TYPE", :install)

    platform_version =
      get_option_value(options[:platform_version],
                       [:puppet5, :puppet6, :puppet7, :puppet8, :"puppet8-nightly"],
                       "Puppet Platform Repo Version",
                       "PLATFORM_VERSION",
                       :puppet5)

    validate_package_version =
        get_option_value(options[:puppetdb_validate_package_version],
            [:true, :false], "'validate package version'",
            "PUPPETDB_VALIDATE_PACKAGE_VERSION", :true)

    expected_rpm_version =
        get_option_value(options[:puppetdb_expected_rpm_version],
            nil, "'expected RPM package version'",
            "PUPPETDB_EXPECTED_RPM_VERSION", nil)

    expected_deb_version =
        get_option_value(options[:puppetdb_expected_deb_version],
                         nil, "'expected DEB package version'",
                         "PUPPETDB_EXPECTED_DEB_VERSION", nil)

    use_proxies =
        get_option_value(options[:puppetdb_use_proxies],
          [:true, :false], "'use proxies'", "PUPPETDB_USE_PROXIES", :false)

    purge_after_run =
        get_option_value(options[:puppetdb_purge_after_run],
          [:true, :false],
          "'purge packages and perform exhaustive cleanup after run'",
          "PUPPETDB_PURGE_AFTER_RUN", :false)

    skip_openvox_package_installation =
        get_option_value(options[:puppetdb_skip_openvox_package_installation],
          [:true, :false],
          "'skip installation of openvox packages, since the test suite does not know how'",
          "PUPPETDB_SKIP_OPENVOX_PACKAGE_INSTALLATION", :true)

    skip_presuite_provisioning =
        get_option_value(options[:puppetdb_skip_presuite_provisioning],
          [:true, :false],
          "'skip installation steps",
          "PUPPETDB_SKIP_PRESUITE_PROVISIONING", :false)

    package_build_host =
        get_option_value(options[:puppetdb_package_build_host],
          nil,
          "'hostname for package build output'",
          "PUPPETDB_PACKAGE_BUILD_HOST",
          "builds.delivery.puppetlabs.net")

    package_repo_host =
        get_option_value(options[:puppetdb_package_repo_host],
          nil,
          "'hostname for yum/apt repos'",
          "PUPPETDB_PACKAGE_REPO_HOST",
          "builds.delivery.puppetlabs.net")

    package_repo_url =
        get_option_value(options[:puppetdb_package_repo_url],
          nil,
          "'base URL for yum/apt repos'",
          "PUPPETDB_PACKAGE_REPO_URL",
          "http://#{package_repo_host}/dev/puppetdb/master")

    package_build_version =
      get_option_value(options[:puppetdb_package_build_version],
        nil, "'package build version'",
        "PUPPETDB_PACKAGE_BUILD_VERSION", nil)

    puppetdb_repo_puppet = get_option_value(options[:puppetdb_repo_puppet],
      nil, "git repo for puppet source installs", "PUPPETDB_REPO_PUPPET", nil)

    puppetdb_repo_hiera = get_option_value(options[:puppetdb_repo_hiera],
      nil, "git repo for hiera source installs", "PUPPETDB_REPO_HIERA", nil)

    puppetdb_repo_facter = get_option_value(options[:puppetdb_repo_facter],
      nil, "git repo for facter source installs", "PUPPETDB_REPO_FACTER", nil)

    puppetdb_git_ref = get_option_value(options[:puppetdb_git_ref],
      nil, "git revision of puppetdb to test against", "REF", nil)

    nightly = get_option_value(options[:nightly], [:true, :false],
      "source puppet-agent and puppetserver from nightly builds", "NIGHTLY", :false)

    @config = {
      :base_dir => base_dir,
      :acceptance_data_dir => File.join(base_dir, "acceptance", "data"),
      :os_families => os_families,
      :install_type => install_type,
      :install_mode => install_mode,
      :platform_version => platform_version,
      :validate_package_version => validate_package_version == :true,
      :expected_rpm_version => expected_rpm_version,
      :expected_deb_version => expected_deb_version,
      :use_proxies => use_proxies == :true,
      :purge_after_run => purge_after_run == :true,
      :package_build_host => package_build_host,
      :package_repo_host => package_repo_host,
      :package_repo_url => package_repo_url,
      :package_build_version => package_build_version,
      :repo_puppet => puppetdb_repo_puppet,
      :repo_hiera => puppetdb_repo_hiera,
      :repo_facter => puppetdb_repo_facter,
      :git_ref => puppetdb_git_ref,
      :skip_openvox_package_installation => skip_openvox_package_installation == :true,
      :skip_presuite_provisioning => skip_presuite_provisioning == :true,
      :nightly => nightly == :true
    }

    pp_config = PP.pp(@config, "")

    Beaker::Log.notify "PuppetDB Acceptance Configuration:\n\n#{pp_config}\n\n"
  end

  class << self
    attr_reader :config
  end

  def self.get_option_value(value, legal_values, description,
    env_var_name = nil, default_value = nil)

    # we give precedence to any value explicitly specified in an options file,
    #  but we also allow environment variables to be used for
    #  puppetdb-specific settings
    value = (value || (env_var_name && ENV[env_var_name]) || default_value)
    if value
      value = value.to_sym
    end

    unless legal_values.nil? or legal_values.include?(value)
      raise ArgumentError, "Unsupported #{description} '#{value}'"
    end

    value
  end

  # Return the configuration hash initialized at the start with
  # initialize_test_config
  #
  # @return [Hash] configuration hash
  def test_config
    PuppetDBExtensions.config
  end

  # Return the fact set for a given hostname
  #
  # Relies on populate_facts to be ran first.
  #
  # @param host [String] hostname to retrieve facts for
  # @return [Hash] facts hash
  def facts(host)
    test_config[:facts][host]
  end

  # Populate the facts storage area of test_config
  #
  # @return [void]
  def populate_facts
    fact_data = hosts.inject({}) do |result, host|
      facts_raw = on host, "facter -y"
      facts = YAML.load(facts_raw.stdout)
      result[host.name] = facts
      result
    end

    test_config[:facts] = fact_data
    nil
  end

  def get_os_family(host)
    result = on(host, "which yum", :silent => true)
    if result.exit_code == 0
      :redhat
    else
      :debian
    end
  end

  def aio_pathing_exists?(host)
    result = on host, %Q|if [ -e /etc/puppetlabs/puppetdb ]; then echo aio; fi|
    result.stdout.strip == "aio"
  end

  def puppetdb_confdir(host, legacy=false)
    if aio_pathing_exists?(host)
      "/etc/puppetlabs/puppetdb"
    else
      "/etc/puppetdb"
    end
  end

  def puppetdb_bin_dir(host)
    "/opt/puppetlabs/server/bin"
  end

  def puppetdb_log_dir(host)
    "/var/log/puppetlabs/puppetdb"
  end

  def puppetdb_pids(host)
    java_bin = "java"
    jar_file = "puppetdb.jar"
    result = on host, %Q(ps -ef | grep "#{java_bin}" | grep "#{jar_file}" | grep " services -c " | awk '{print $2}')
    pids = result.stdout.chomp.split("\n")
    Beaker::Log.notify "PuppetDB PIDs appear to be: '#{pids}'"
    pids
  end

  def start_puppetdb(host)
    step "Starting PuppetDB" do
      on(host, puppet_resource('service', 'puppetdb', 'ensure=running'))
      sleep_until_started(host)
    end
  end

  def check_record_count(endpoint, expected_count)
    result = on database, %Q|curl -G http://localhost:8080/pdb/query/v4/#{endpoint}|
    result_records = JSON.parse(result.stdout)
    assert_equal(expected_count, result_records.length,
                 "Expected query to return '#{expected_count}' records from /pdb/query/v4/#{endpoint}; returned '#{result_records.length}'")
  end

  # Display the last few log lines for debugging purposes
  #
  # @param host Hostname to display results for
  # @return [void]
  # @api private
  def display_last_logs(host)
    on host, "tail -n 100 #{puppetdb_log_dir(host)}/puppetdb-daemon.log", :acceptable_exit_codes => [0,1]
    on host, "tail -n 100 #{puppetdb_log_dir(host)}/puppetdb.log", :acceptable_exit_codes => [0,1]
  end

  # Sleep until PuppetDB is completely started
  #
  # @param host Hostname to test for PuppetDB availability
  # @return [void]
  # @api public
  def sleep_until_started(host)
    # Hit an actual endpoint to ensure PuppetDB is up and not just the webserver.
    # Retry until an HTTP response code of 200 is received.
    test_route = aio_pathing_exists?(host) ? "pdb/meta/v1/version" : "v4/version"
    curl_with_retries("start puppetdb", host,
                      "-s -w '%{http_code}' http://localhost:8080/#{test_route} -o /dev/null",
                      0, 120, 1, /200/)
    curl_with_retries("start puppetdb (ssl)", host,
                      "https://#{host.node_name}:8081/#{test_route}", [35, 60])
  rescue RuntimeError => _e
    display_last_logs(host)
    raise
  end

  def is_bionic()
    return test_config[:os_families].has_key? 'ubuntu1804-64-1'
  end

  def is_focal()
    return test_config[:os_families].has_key? 'ubuntu2004-64-1'
  end

  def is_jammy()
    return test_config[:os_families].has_key? 'ubuntu2204-64-1'
  end

  def is_buster()
    return test_config[:os_families].has_key? 'debian10-64-1'
  end

  def is_bullseye()
    return test_config[:os_families].has_key? 'debian11-64-1'
  end

  def is_bookworm()
    return test_config[:os_families].has_key? 'debian12-64-1'
  end

  def is_el8()
    return test_config[:os_families].has_key?('redhat8-64-1') ||
           test_config[:os_families].has_key?('centos8-64-1')
  end

  def is_el9()
    return test_config[:os_families].has_key?('redhat9-64-1')
  end

  def is_rhel7fips
    return test_config[:os_families].has_key? 'redhatfips7-64-1'
  end

  # This specifies which puppet platform repo to initialize in order
  # to find packages.
  #
  # For install tests it should pick the nightly repo so we are testing the
  # latest dev builds of agent and server
  #
  # For upgrade_latest tests it should be the non-nightly version. Tests are
  # run with the shipped package because nightly packages are pruned eventually.
  #
  # For upgrade_oldest tests it must match the result of oldest_supported
  #
  # TODO: when testing X.0.0-SNAPSHOT one of install/upgrade_latest tests
  # will fail because those two tests will need to use different repos
  def puppet_repo_version(platform_version, install_type, nightly)
    # When adding a new platform version, only puppet8-nightly will be available
    # for a time. In order to generate packages we will have to run with the nightly
    # platform version (in addition to our own nightly boolean). This will prune the -nightly
    # and allow us to function normally.
    sanitized_version = "#{platform_version}".delete_suffix("-nightly").to_sym

    case install_type
    when :install
      if nightly
        :"#{sanitized_version}-nightly"
      else
        sanitized_version
      end
    when :upgrade_latest
      # always use the non-nightly version to install released packages
      sanitized_version
    when :upgrade_oldest
      if is_bookworm
        :puppet8
      else
        :puppet7
      end
    end
  end

  # This accounts for packages we added _after_ we released the oldest
  # supported version of PuppetDB. Its version must be available in the
  # platform version returned by puppet_repo_version above
  def oldest_supported
    # account for bionic/rhel8 not having build before certain versions
    if is_el9 || is_jammy
      '7.14.0'
    elsif is_bullseye
      '7.9.0'
    elsif is_bookworm
      '8.8.0'
    else
      '7.3.1'
    end
  end

  def get_latest_tag(branch)
    Open3.capture2('git', '-C', 'puppetdb', 'describe', '--tags', '--abbrev=0', "origin/#{branch}")
  end

  def get_latest_released(version)
    cloned = system('git clone https://github.com/puppetlabs/puppetdb.git')
    if cloned.nil?
      raise 'error cloning puppetdb repo'
    end

    # infer branch from version, note this will ask for a non existent branch when on master
    branch = /^(\d+\.)\d+\.\d+/.match(version)[1] + 'x'
    tag, status = get_latest_tag(branch)

    if status.success?
      return tag.delete("\n")
    else
      # if tag doesn't exist assume branch is main
      tag, status = get_latest_tag('main')
      if status.success?
        return tag.delete("\n")
      else
        raise "error getting most recent tagged release. status: #{status}"
      end
    end
  end

  def get_package_version(host, version = nil)
    if version == 'latest'
      return 'latest'
    elsif version.nil?
      version = PuppetDBExtensions.config[:package_build_version].to_s
      # If no version was defined, default to latest.
      if version == ''
        return 'latest'
      end
    end

    # version can look like:
    #   3.0.0
    #   3.0.0.SNAPSHOT.2015.07.08T0945

    # Rewrite version if its a SNAPSHOT in rc form
    if version.include?("SNAPSHOT")
      version = version.sub(/^(.*)\.(SNAPSHOT\..*)$/, "\\1-0.1\\2")
    else
      version = version + "-1"
    end

    ## These 'platform' values come from the acceptance config files, so
    ## we're relying entirely on naming conventions here.  Would be nicer
    ## to do this using lsb_release or something, but...
    if host['platform'].include?('el-7')
      "#{version}.el7"
    elsif host['platform'].include?('el-8')
      "#{version}.el8"
    elsif host['platform'].include?('el-9')
      "#{version}.el9"
    elsif host['platform'].include?('ubuntu-18.04')
      "#{version}bionic"
    elsif host['platform'].include?('ubuntu-20.04')
      "#{version}focal"
    elsif host['platform'].include?('ubuntu-22.04')
      "#{version}jammy"
    elsif host['platform'].include?('debian-10')
      "#{version}buster"
    elsif host['platform'].include?('debian-11')
      "#{version}bullseye"
    elsif host['platform'].include?('debian-12')
      "#{version}bookworm"
    else
      raise ArgumentError, "Unsupported platform: '#{host['platform']}'"
    end
  end

  def install_puppetdb_module(hosts, puppet_platform)
    on(hosts, 'puppet module install puppetlabs-puppetdb')
  end

  def setup_openvoxdb_certs(database)
    step 'Ensure PuppetDB certificates are setup' do
      # Normally the openvoxdb package automagically post-install runs
      # 'puppetdb ssl-setup', but this relies on the openvox-agent
      # already having certs generated at the time that the openvoxdb
      # package was installed. If we installed openvoxdb before
      # running the suite and before agent certs were generated, then
      # we need this step to ensure that openvoxdb now gets its certs.
      # (If they are already in place, the command should be
      # idempotent.)
      apply_manifest_on(database, <<~EOM)
        exec { 'puppetdb-prepare-certs':
          command => '/opt/puppetlabs/bin/puppetdb ssl-setup',
          path    => '/bin:/sbin:/usr/bin',
          onlyif  => 'test -f /opt/puppetlabs/bin/puppetdb',
        }
      EOM
    end
  end

  # Use the puppetdb module to configure pre-installed openvoxdb
  # and packages.
  def configure_openvoxdb(host)
    manifest = <<~EOS
      class { 'puppetdb':
        puppetdb_package    => 'openvoxdb',
        manage_firewall     => false,
        manage_package_repo => true,
        disable_update_checking => true,
        database_listen_address => 'localhost',
        postgres_version        => '14',
        database_name           => 'puppetdb',
        database_username       => 'puppetdb',
        database_password       => 'puppetdb',
        read_database_username  => 'puppetdb-read',
        read_database_password  => 'puppetdb-read',
      }
    EOS
    apply_manifest_on(host, manifest)
    print_ini_files(host)
    sleep_until_started(host)
  end

  def configure_openvoxdb_termini(host, databases)
    server_urls = databases.map {|db| "https://#{db.node_name}:8081"}.join(',')
    manifest = <<~EOS
      class { 'puppetdb::master::config':
        terminus_package         => 'openvoxdb-termini',
        puppetdb_startup_timeout => 120,
        manage_report_processor  => true,
        enable_reports           => true,
        strict_validation        => true,
      }
      ini_setting {'server_urls':
        ensure  => present,
        section => 'main',
        path    => "${puppetdb::params::puppet_confdir}/puppetdb.conf",
        setting => 'server_urls',
        value   => '#{server_urls}',
      }
    EOS
    apply_manifest_on(host, manifest)
  end

  def install_puppetdb(host, version=nil)
    manifest = <<-EOS
    class { 'puppetdb::globals':
      version => '#{get_package_version(host, version)}'
    }
    class { 'puppetdb::server':
      manage_firewall => false,
      disable_update_checking => true,
    }
    #{postgres_manifest}
    Class['::puppetdb::database::postgresql'] -> Class['puppetdb::server']
    EOS
    apply_manifest_on(host, manifest)
    print_ini_files(host)
    sleep_until_started(host)
  end

  def validate_package_version(host)
    step "Verifying package version" do
      os = PuppetDBExtensions.config[:os_families][host.name]
      installed_version =
        case os
          when :debian
            result = on host, "dpkg-query --showformat \"\\${Version}\" --show puppetdb"
            result.stdout.strip
          when :redhat
            result = on host, "rpm -q puppetdb --queryformat \"%{VERSION}-%{RELEASE}\""
            result.stdout.strip
          else
            raise ArgumentError, "Unsupported OS family: '#{os}'"
        end
      expected_version = get_package_version(host)

      Beaker::Log.notify "Expecting package version: '#{expected_version}', actual version: '#{installed_version}'"
      if installed_version != expected_version and expected_version != 'latest'
        raise RuntimeError, "Installed version '#{installed_version}' did not match expected version '#{expected_version}'"
      end
    end
  end

  def install_puppetdb_termini(host, databases, version=nil)
    # We pass 'restart_puppet' => false to prevent the module from trying to
    # manage the puppet master service, which isn't actually installed on the
    # acceptance nodes (they run puppet master from the CLI).
    server_urls = databases.map {|db| "https://#{db.node_name}:8081"}.join(',')
    manifest = <<-EOS
    class { 'puppetdb::globals':
      version => '#{get_package_version(host, version)}'
    }
    class { 'puppetdb::master::config':
      puppetdb_startup_timeout => 120,
      manage_report_processor  => true,
      enable_reports           => true,
      strict_validation        => true,
      restart_puppet           => false,
    }
    ini_setting {'server_urls':
      ensure  => present,
      section => 'main',
      path    => "${puppetdb::params::puppet_confdir}/puppetdb.conf",
      setting => 'server_urls',
      value   => '#{server_urls}',
    }
    EOS
    apply_manifest_on(host, manifest)
  end

  def print_ini_files(host)
    confdir = "#{puppetdb_confdir(host)}/conf.d"
    step "Print out jetty.ini for posterity" do
      on host, "cat #{confdir}/jetty.ini"
    end
    step "Print out database.ini for posterity" do
      on host, "cat #{confdir}/database.ini"
    end
  end

  def current_time_on(host)
    result = on host, %Q|date --rfc-2822|
    CGI.escape(Time.rfc2822(result.stdout).iso8601)
  end

  def enable_https_apt_sources(host)
    manifest = <<-EOS
      if ($facts['os']['family'] == 'Debian') {
        package { 'apt-transport-https': ensure => installed }
      }
    EOS

    apply_manifest_on(host, manifest)
  end

  # Work around for testing on rhel and the repos on the image
  # not finding the pg packages it needs.
  # TODO: figure out why puppetdb module isn't managing this?
  def configure_postgresql_repos_on_el(host)
    if host.platform =~ /el-/
      step 'Update EL postgresql repos' do
        major_version = host.platform.split('-')[1]

        on(host, "dnf install -y https://download.postgresql.org/pub/repos/yum/reporpms/EL-#{major_version}-x86_64/pgdg-redhat-repo-latest.noarch.rpm")
        on(host, "dnf -qy module disable postgresql")
      end
    end
  end

  def postgres_manifest
    # For EOL OSes, pgdg removes the package repo
    manage_package_repo = ! (is_bionic || is_buster)

    manifest = <<~EOS
      # create the puppetdb database
      class { '::puppetdb::database::postgresql':
        listen_addresses            => 'localhost',
        manage_package_repo         => #{manage_package_repo},
        postgres_version            => '14',
        database_name               => 'puppetdb',
        database_username           => 'puppetdb',
        database_password           => 'puppetdb',
        read_database_username      => 'puppetdb-read',
        read_database_password      => 'puppetdb-read',
      }
    EOS
    manifest
  end

  def install_postgres(host)
    Beaker::Log.notify "Installing postgres on #{host}"
    apply_manifest_on(host, postgres_manifest)
  end

  # Restart postgresql using Puppet, by notifying the postgresql::server::service
  # class, which should cause the service to restart.
  def restart_postgres(host)
    manifest = <<-EOS
      notify { 'restarting postgresql': }~>
      Class['postgresql::server::service']
      #{postgres_manifest}
    EOS
    apply_manifest_on(host, manifest)
  end

  # Prepare host's global git configuration so git tag will just work
  # later on.
  #
  # @param host [String] host to execute on
  # @return [void]
  # @api private
  def prepare_git_author(host)
    on host, 'git config --global user.email "beaker@localhost"'
    on host, 'git config --global user.name "Beaker"'
  end

  # Install puppetdb using `rake install` methodology, which works
  # without packages today.
  #
  # @param host [String] host to execute on
  # @return [void]
  # @api public
  def install_puppetdb_via_rake(host)
    # Uncomment for pinning against a particular ezbake revision
    #ezbake_dev_build("git@github.com:kbarber/ezbake.git", "pdb-1455")

    install_from_ezbake host

    step "Configure database.ini file" do
      manifest = "
        class { 'puppetdb::server::database_ini': }"

      apply_manifest_on(host, manifest)
    end

    print_ini_files(host)
  end

  def install_puppetdb_termini_via_rake(host, databases)
    # Uncomment for pinning against a particular ezbake revision
    #ezbake_dev_build("git@github.com:kbarber/ezbake.git", "pdb-1455")

    install_termini_from_ezbake host

    server_urls = databases.map {|db| "https://#{db.node_name}:8081"}.join(',')
    manifest = <<-EOS
      class { 'puppetdb::globals':
        version => '#{get_package_version(host)}'
      }
      include puppetdb::master::routes
      include puppetdb::master::storeconfigs
      class { 'puppetdb::master::report_processor':
        enable => true,
      }
      ini_setting {'server_urls':
        ensure => present,
        section => 'main',
        path => "${puppetdb::params::puppet_confdir}/puppetdb.conf",
        setting => 'server_urls',
        value => '#{server_urls}',
      }
    EOS
    apply_manifest_on(host, manifest)
  end

  def stop_puppetdb(host)
    step 'Stopping PuppetDB' do
      pids = puppetdb_pids(host)
      on(host, puppet_resource('service', 'puppetdb', 'ensure=stopped'))
      sleep_until_stopped(host, pids)
    end
  end

  def sleep_until_stopped(host, pids)
    curl_with_retries("stop puppetdb", host, "http://localhost:8080", 7)
    stopped = false
    while (! stopped)
      Beaker::Log.notify("Waiting for pids #{pids} to exit")
      exit_codes = pids.map do |pid|
        result = on host, %Q(ps -p #{pid}), :acceptable_exit_codes => [0, 1]
        result.exit_code
      end
      stopped = exit_codes.all? { |x| x == 1}
    end
  end

  def restart_puppetdb(host)
    stop_puppetdb(host)
    start_puppetdb(host)
  end

  def clear_and_restart_puppetdb(host)
    stop_puppetdb(host)
    clear_database(host)
    start_puppetdb(host)
  end

  def sleep_until_queue_empty(host, timeout=60)
    metric = "puppetlabs.puppetdb.mq:name=global.depth"
    queue_size = nil

    begin
      Timeout.timeout(timeout) do
        until queue_size == 0
          result = on host, %Q(curl http://localhost:8080/metrics/v1/mbeans/#{CGI.escape(metric)} 2> /dev/null | python -c "import sys, json; print json.load(sys.stdin)['Count']")
          queue_size = Integer(result.stdout.chomp)
        end
      end
    rescue Timeout::Error => _e
      raise "Queue took longer than allowed #{timeout} seconds to empty"
    end
  end

  # Queries the metrics endpoint for command processing results, return a hash
  # of results.
  def command_processing_stats(host, counter = "processed")
    metric = "puppetlabs.puppetdb.command:type=global,name=discarded"

    result = on host, %Q(curl http://localhost:8080/metrics/v1/mbeans/#{CGI.escape(metric)})

    begin
      JSON.parse(result.stdout)
    rescue JSON::ParserError => e
      Beaker::Log.notify "Invalid JSON response: #{result.stdout}"
      raise e
    end
  end

  def apply_manifest_on(host, manifest_content)
    manifest_path = host.tmpfile("puppetdb_manifest.pp")
    create_remote_file(host, manifest_path, manifest_content)
    Beaker::Log.notify "Applying manifest on #{host}:\n\n#{manifest_content}"
    on host, puppet_apply("--detailed-exitcodes #{manifest_path}"), :acceptable_exit_codes => [0,2]
  end

  def modify_config_setting(host, config_file_name, section, setting, value)
    manifest_content = <<EOS
ini_setting {'puppetdb-config-#{setting}':
    path => '#{puppetdb_confdir(host)}/conf.d/#{config_file_name}',
    section => '#{section}',
    setting => '#{setting}',
    value => '#{value}'
}
EOS
    apply_manifest_on(host, manifest_content)
  end

  # Keep curling until the required condition is met
  #
  # Condition can be a desired_exit code, and/or expected_output, and it will
  # keep executing the curl command until one of these conditions are met
  # or the max_retries is reached.
  #
  # @param desc [String] descriptive name for this cycling
  # @param host [String] host to execute retry on
  # @param curl_args [String] the curl_args to use for testing
  # @param desired_exit_codes [Number,Array<Number>] a desired exit code, or array of exist codes
  # @param max_retries [Number] maximum number of retries before failing
  # @param retry_interval [Number] time in secs to wait before next try
  # @param expected_output [Regexp] a regexp to use for matching the output
  # @return [void]
  def curl_with_retries(desc, host, curl_args, desired_exit_codes, max_retries = 60, retry_interval = 1, expected_output = /.*/)
    command = "curl --tlsv1.3 #{curl_args}"
    log_prefix = host.log_prefix
    logger.debug "\n#{log_prefix} #{Time.new.strftime('%H:%M:%S')}$ #{command}"
    logger.debug "  Trying command #{max_retries} times."
    logger.debug ".", add_newline=false

    desired_exit_codes = [desired_exit_codes].flatten
    result = on host, command, :acceptable_exit_codes => (0...127), :silent => true
    num_retries = 0
    until desired_exit_codes.include?(result.exit_code) and (result.stdout =~ expected_output)
      sleep retry_interval
      result = on host, command, :acceptable_exit_codes => (0...127), :silent => true
      num_retries += 1
      logger.debug ".", add_newline=false
      if (num_retries > max_retries)
        logger.debug "  Command \`#{command}\` failed."
        fail("Command \`#{command}\` failed. Unable to #{desc}.")
      end
    end
    logger.debug "\n#{log_prefix} #{Time.new.strftime('%H:%M:%S')}$ #{command} ostensibly successful."
  end

  def clear_database(host)
    on host, 'su postgres -c "dropdb puppetdb"'
    install_postgres(host)
  end

  # Restore puppet.conf from backup, if puppet.conf.bak exists.
  #
  # @api private
  def restore_puppet_conf host
    confdir = host.puppet['confdir']
    on host, "if [ -f #{confdir}/puppet.conf.bak ]; then " +
               "cat #{confdir}/puppet.conf.bak > " +
               "#{confdir}/puppet.conf; " +
               "rm -rf #{confdir}/puppet.conf.bak; " +
             "fi"
  end

  def parse_json_with_error(input)
    begin
      facts = JSON.parse(input)
    rescue Exception => e
      facts = "#{e.message} on input '#{input}'"
    end
    return facts
  end

  ##############################################################################
  # Object diff functions
  ##############################################################################
  # This is horrible and really doesn't belong here, but I'm not sure where
  # else to put it.  I need a way to do a recursive diff of a hash (which may
  # contain nested objects whose type can be any of Hash, Array, Set, or a
  # scalar).  The hashes may be absolutely gigantic, so if they don't match,
  # I need a way to be able to show a small enough diff so that the user can
  # actually figure out what's going wrong (rather than dumping out the entire
  # gigantic string).  I searched for gems that could handle this and tried
  # 4 or 5 different things, and couldn't find anything that suited the task,
  # so I had to write my own.  This could use improvement, relocation, or
  # replacement with a gem if we ever find a suitable one.
  #
  # UPDATE: chatted with Justin about this and he suggests creating a special
  # puppetlabs-test-utils repo or similar and have that pulled down via
  # bundler, once the acceptance harness is accessible as a gem.  You know,
  # in "The Future".

  # JSON gem doesn't have native support for Set objects, so we have to
  # add this hack.
  class ::Set
    def to_json(arg)
      to_a.to_json(arg)
    end
  end

  def hash_diff(obj1, obj2)
    result =
      (obj1.keys | obj2.keys).inject({}) do |diff, k|
        if obj1[k] != obj2[k]
          objdiff = object_diff(obj1[k], obj2[k])
          if (objdiff)
            diff[k] = objdiff
          end
        end
        diff
      end
    (result == {}) ? nil : result
  end

  def array_diff(arr1, arr2)
    (0..([arr1.length, arr2.length].max)).inject([]) do |diff, i|
      objdiff = object_diff(arr1[i], arr2[i])
      if (objdiff)
        diff << objdiff
      end
      diff
    end
  end

  def set_diff(set1, set2)
    diff1 = set1 - set2
    diff2 = set2 - set1
    unless (diff1.empty? and diff2.empty?)
      [diff1, diff2]
    end
  end

  def object_diff(obj1, obj2)
    if (obj1.class != obj2.class)
      [obj1, obj2]
    else
      case obj1
        when Hash
          hash_diff(obj1, obj2)
        when Array
          array_diff(obj1, obj2)
        when Set
          set_diff(obj1, obj2)
        else
          (obj1 == obj2) ? nil : [obj1, obj2]
      end
    end
  end

  ##############################################################################
  # End Object diff functions
  ##############################################################################

  def install_puppet_from_package(puppet_collection = "puppet6")
    hosts.each do |host|
      install_puppet_agent_on(host, {:puppet_collection => puppet_collection})
      on( host, puppet('resource', 'host', 'updates.puppetlabs.com', 'ensure=present', "ip=127.0.0.1") )

      # Take new configuration file on debian to prevent tests from hanging
      # when this is used to upgrade puppet
      if get_os_family(database) == :debian
        # Explicitly install ca-certificates-java before java to avoid error installing java on debian
        database.install_package('ca-certificates-java')
        database.install_package('puppetserver', '-o Dpkg::Options::="--force-confnew"')
      else
        install_package(database, 'puppetserver')
      end
    end
  end

  # This helper has been grabbed from beaker, and overriden with the opts
  # component so I can add a new 'refspec' functionality to allow a custom
  # refspec if required.
  #
  # Once this methodology is confirmed we should merge it back upstream.
  def install_from_git(host, path, repository, opts = {})
    name   = repository[:name]
    repo   = repository[:path]
    rev    = repository[:rev]

    target = "#{path}/#{name}"

    step "Clone #{repo} if needed" do
      on host, "test -d #{path} || mkdir -p #{path}"
      on host, "test -d #{target} || git clone #{repo} #{target}"
    end

    step "Update #{name} and check out revision #{rev}" do
      commands = ["cd #{target}",
                  "remote rm origin",
                  "remote add origin #{repo}",
                  "fetch origin #{opts[:refspec]}",
                  "clean -fdx",
                  "checkout -f #{rev}"]
      on host, commands.join(" && git ")
    end

    step "Install #{name} on the system" do
      # The solaris ruby IPS package has bindir set to /usr/ruby/1.8/bin.
      # However, this is not the path to which we want to deliver our
      # binaries. So if we are using solaris, we have to pass the bin and
      # sbin directories to the install.rb
      install_opts = ''
      install_opts = '--bindir=/usr/bin --sbindir=/usr/sbin' if
        host['platform'].include? 'solaris'

        on host,  "cd #{target} && " +
                  "if [ -f install.rb ]; then " +
                  "ruby ./install.rb #{install_opts}; " +
                  "else true; fi"
    end
  end

  def install_puppet_from_source
    os_families = test_config[:os_families]

    extend Beaker::DSL::InstallUtils

    source_path = Beaker::DSL::InstallUtils::SourcePath
    git_uri     = Beaker::DSL::InstallUtils::GitURI
    github_sig  = Beaker::DSL::InstallUtils::GitHubSig

    tmp_repositories = []

    repos = Hash[*test_config.select {|k, v| k =~ /^repo_/}.flatten].values

    repos.each do |uri|
      raise(ArgumentError, "#{uri} is not recognized.") unless(uri =~ git_uri)
      tmp_repositories << extract_repo_info_from(uri)
    end

    repositories = order_packages(tmp_repositories)

    hosts.each_with_index do |host, index|
      os = os_families[host.name]

      case os
      when :redhat
        install_pacakge(host, 'ruby')
        install_pacakge(host, 'git-core')
      when :debian
        install_pacakge(host, 'ruby')
        install_pacakge(host, 'git-core')
      else
        raise "OS #{os} not supported"
      end

      on host, "echo #{github_sig} >> $HOME/.ssh/known_hosts"

      repositories.each do |repository|
        step "Install #{repository[:name]}"
        install_from_git host, source_path, repository,
          :refspec => '+refs/pull/*:refs/remotes/origin/pr/*'
      end

      on host, "getent group puppet || groupadd puppet"
      on host, "getent passwd puppet || useradd puppet -g puppet -G puppet"
      on host, "mkdir -p /var/run/puppet"
      on host, "chown puppet:puppet /var/run/puppet"
    end
  end

  def install_puppet_conf
    hosts.each do |host|
      hostname = fact_on(host, "hostname")
      fqdn = fact_on(host, "fqdn")

      confdir = host.puppet['confdir']
      on host, "mkdir -p #{confdir}"


      puppetconf = File.join(confdir, 'puppet.conf')
      pidfile = '/var/run/puppetlabs/puppetserver/puppetserver.pid'
      conf = IniFile.new
      conf['agent']['server'] = master
      conf['master']['pidfile'] = pidfile
      conf['master']['dns_alt_names']="puppet,#{hostname},#{fqdn},#{host.hostname}"
      create_remote_file(host, puppetconf, conf.to_s)
    end
  end

  def install_puppet(puppet_collection = "puppet6")
    # If our :install_type is :pe then the harness has already installed puppet.
    case test_config[:install_type]
    when :package
      install_puppet_from_package(puppet_collection)
    when :git
      if test_config[:repo_puppet] then
        install_puppet_from_source
      else
        puts "Error: No puppet repo set for installation from git source"
        exit 1
      end
    else
      puts "Error: No puppet repo set for installation from git source"
      exit 1
    end
    install_puppet_conf
  end

  def databases
    extend Beaker::DSL::Roles
    hosts_as(:database).sort_by {|db| db.to_str}
  end

  def database
    # primary database must be numbered lowest
    databases[0]
  end
end

module Beaker
  module DSL
    module Helpers
      module PuppetServerAcceptance
        # A sad little abstraction around service/systemctl to handle
        # os differences for edge cases not suited to puppet resource
        # service calls.
        def service(host, action, service_name, acceptable_exit_codes: [0])
          if find_executable('systemctl')
            command = Command.new("systemctl #{action} #{service_name}")
          elsif find_executable('service')
            command = Command.new("service #{service_name} #{action}")
          else
            raise "Neither systemctl nor service found on #{host.name}"
          end

          host.exec(command, acceptable_exit_codes: acceptable_exit_codes)
        end

        # Override beaker-puppet BeakerPuppet::Helpers::PuppetHelpers#bounce_service
        # to allow for systems that only have systemctl now (el9+, etc.)
        #
        # Ostensibly we should fork beaker-puppet...I'm just not quite
        # ready to touch that yet...
        #
        # Restarts the named puppet service
        #
        # @param [Host] host Host the service runs on
        # @param [String] service Name of the service to restart
        # @param [Fixnum] curl_retries Number of seconds to wait for the restart to complete before failing
        # @param [Fixnum] port Port to check status at
        #
        # @return [Result] Result of last status check
        # @!visibility private
        def bounce_service(host, service, curl_retries = nil, port = nil)
          curl_retries = 120 if curl_retries.nil?
          port = options[:puppetserver_port] if port.nil?
          result = service(host, :reload, service, acceptable_exit_codes: [0, 1, 3])
          return result if result.exit_code == 0

          host.exec puppet_resource('service', service, 'ensure=stopped')
          host.exec puppet_resource('service', service, 'ensure=running')

          curl_with_retries(" #{service} ", host, "https://localhost:#{port}", [35, 60], curl_retries)
        end
      end
    end
  end
end

# Register the DSL extension
Beaker::DSL.register(Beaker::DSL::Helpers::PuppetServerAcceptance)

# oh dear.
Beaker::TestCase.send(:include, PuppetDBExtensions)
