require 'rake/clean'
require 'pp'
require 'yaml'
require 'securerandom'
require 'fileutils'
require 'beaker-hostgenerator'
require 'beaker/dsl/install_utils'
extend Beaker::DSL::InstallUtils
$LOAD_PATH << File.expand_path(File.join(File.dirname(__FILE__), 'lib'))

ONE_DAY_IN_SECS = 24 * 60 * 60
REPO_CONFIGS_DIR = "repo-configs"
CLEAN.include('*.tar', REPO_CONFIGS_DIR, 'merged_options.rb')

# If elsewhere we're depending on internal network resources
# then assume we can depend on them here
EPEL_MIRROR = ENV['BEAKER_EPEL_MIRROR']

# Default test target if none specified
DEFAULT_MASTER_TEST_TARGET = 'redhat7-64m'
DEFAULT_TEST_TARGETS = "#{DEFAULT_MASTER_TEST_TARGET}a-windows2012r2-64a"

module HarnessOptions
  defaults = {
    :type => 'git',
    :helper => ['lib/helper.rb'],
    :tests  => ['tests'],
    :log_level => 'debug',
    :color => false,
    :root_keys => true,
    :ssh => {
      :keys => ["id_rsa_acceptance", "#{ENV['HOME']}/.ssh/id_rsa-acceptance"],
    },
    :xml => true,
    :timesync => false,
    :repo_proxy => true,
    :add_el_extras => false,
    :preserve_hosts => 'onfail',
    :forge_host => 'forge-aio01-petest.puppetlabs.com',
    :'master-start-curl-retries' => 30
  }

  DEFAULTS = EPEL_MIRROR ? defaults.merge(:epel_url => EPEL_MIRROR) : defaults

  class Aggregator
    attr_reader :mode

    def initialize(mode)
      @mode = mode
    end

    def get_options(file_path)
      puts file_path
      if File.exists? file_path
        options = eval(File.read(file_path), binding)
      else
        puts "No options file found at #{File.expand_path(file_path)}"
      end
      options || {}
    end

    def get_mode_options
      get_options("./config/#{mode}/options.rb")
    end

    def get_local_options
      get_options("./local_options.rb")
    end

    def final_options(intermediary_options = {})
      mode_options = get_mode_options
      local_overrides = get_local_options
      final_options = DEFAULTS.merge(mode_options)
      final_options.merge!(intermediary_options)
      final_options.merge!(local_overrides)
      return final_options
    end
  end

  def self.options(mode, options)
    final_options = Aggregator.new(mode).final_options(options)
    final_options
  end
end

def beaker_test(mode = :aio, options = {})
  delete_options = options.delete(:__delete_options__) || []
  final_options = HarnessOptions.options(mode, options)
  preserve_config = final_options.delete(:__preserve_config__)

  if mode == :git
    # Build up project git urls based on git server and fork env variables or defaults
    final_options[:install].map! do |install|
      raise(ArgumentError, "Missing Git URL within options hash. Install URL is nil.") if install.nil?
      if md = /^(\w+)#([\d.\w]+)$/.match(install)
        project, project_sha = md.captures
        "#{build_giturl(project)}##{project_sha}"
      elsif md = /^(\w+)$/.match(install)
        project = md[1]
        "#{build_giturl(project)}##{sha}"
      else
        install
      end
    end
  end

  delete_options.each do |delete_me|
    final_options.delete(delete_me)
  end

  options_file = 'merged_options.rb'
  File.open(options_file, 'w') do |merged|
    merged.puts <<-EOS
# Copy this file to local_options.rb and adjust as needed if you wish to run
# with some local overrides.
EOS
    merged.puts(final_options.pretty_inspect)
  end

  tests = ENV['TESTS'] || ENV['TEST']
  tests_opt = "--tests=#{tests}" if tests

  if config and File.exists?(config)
    config_opt = "--hosts=#{config}"
  else
    if agent_target = ENV['TEST_TARGET']
      master_target = ENV['MASTER_TEST_TARGET'] || DEFAULT_MASTER_TEST_TARGET
      targets = "#{master_target}-#{agent_target}"
    else
      targets = DEFAULT_TEST_TARGETS
    end
    cli = BeakerHostGenerator::CLI.new([targets, '--disable-default-role', '--osinfo-version', '1'])
    ENV['BEAKER_HOSTS'] = "tmp/#{targets}-#{SecureRandom.uuid}.yaml"
    FileUtils.mkdir_p('tmp')
    File.open(config, 'w') do |fh|
      fh.print(cli.execute)
    end
    config_opt = "--hosts=#{config}"
  end

  overriding_options = ENV['OPTIONS']

  args = ["--options-file", options_file, config_opt, tests_opt, overriding_options].compact

  begin
    sh("beaker", *args)
  ensure
    preserve_configuration(final_options, options_file) if preserve_config
  end
end

def preserve_configuration(final_options, options_file)
  if (hosts_file = config || final_options[:hosts_file]) && hosts_file !~ /preserved_config/
    cp(hosts_file, "log/latest/config.yml")
    generate_config_for_latest_hosts
  end
  mv(options_file, "log/latest")
end

def generate_config_for_latest_hosts
  preserved_config_hash = { 'HOSTS' => {} }

  puts "\nPreserving configuration so that any preserved nodes can be tested again locally..."

  config_hash = YAML.load_file('log/latest/config.yml')
  if !config_hash || !config_hash.include?('HOSTS')
    puts "Warning: No HOSTS configuration found in log/latest/config.yml"
    return
  else
    nodes = config_hash['HOSTS'].map do |node_label,hash|
      {
        :node_label => node_label,
        :roles => hash['roles'],
        :platform => hash['platform']
      }
    end

    pre_suite_log = File.read('log/latest/pre_suite-run.log')
    nodes.each do |node_info|
      host_regex = /^([\w.]+) \(#{node_info[:node_label]}\)/
      if matched = host_regex.match(pre_suite_log)
        hostname = matched[1]
        fqdn = (hostname =~ /\./) ?
          hostname :
          "#{hostname}.delivery.puppetlabs.net"
      elsif /^#{node_info[:node_label]} /.match(pre_suite_log)
        fqdn = "#{node_info[:node_label]}"
        puts "* Couldn't find any log lines for #{host_regex}, assuming #{fqdn} is the fqdn"
      end
      if fqdn
        preserved_config_hash['HOSTS'][fqdn] = {
          'roles' => node_info[:roles],
          'platform' => node_info[:platform],
        }
      else
        puts "* Couldn't match #{node_info[:node_label]} in pre_suite-run.log"
      end
    end
    pp preserved_config_hash

    File.open('log/latest/preserved_config.yaml', 'w') do |config_file|
      YAML.dump(preserved_config_hash, config_file)
    end
  end
rescue Errno::ENOENT => e
  puts "Warning: Couldn't generate preserved_config.yaml #{e}"
end

def list_preserved_configurations(secs_ago = ONE_DAY_IN_SECS)
  preserved = {}
  Dir.glob('log/*_*').each do |dir|
    preserved_config_path = "#{dir}/preserved_config.yaml"
    yesterday = Time.now - secs_ago.to_i
    if preserved_config = File.exists?(preserved_config_path)
      directory = File.new(dir)
      if directory.ctime > yesterday
        hosts = []
        preserved_config = YAML.load_file(preserved_config_path).to_hash
        preserved_config['HOSTS'].each do |hostname,values|
          hosts << "#{hostname}: #{values['platform']}, #{values['roles']}"
        end
        preserved[hosts] = directory.to_path
      end
    end
  end
  preserved.map { |k,v| [v,k] }.sort { |a,b| a[0] <=> b[0] }.reverse
end

def list_preserved_hosts(secs_ago = ONE_DAY_IN_SECS)
  hosts = Set.new
  Dir.glob('log/**/pre*suite*run.log').each do |log|
    yesterday = Time.now - secs_ago.to_i
    File.open(log, 'r') do |file|
      if file.ctime > yesterday
        file.each_line do |line|
          matchdata = /^(\w+)(?:\.[\w.]+)? \(.*?\) \d\d:\d\d:\d\d\$/.match(line.encode!('UTF-8', 'UTF-8', :invalid => :replace))
          hosts.add(matchdata[1]) if matchdata
        end
      end
    end
  end
  hosts
end

def release_hosts(hosts = nil, secs_ago = ONE_DAY_IN_SECS)
  secs_ago ||= ONE_DAY_IN_SECS
  hosts ||= list_preserved_hosts(secs_ago)

  hosts.each do |h|
    hostname = h.split('.').first
    puts "Releaseing '#{hostname}'"
    puts `curl -X DELETE --url http://vcloud.delivery.puppetlabs.net/vm/#{hostname}`
  end
end

def print_preserved(preserved)
  preserved.each_with_index do |entry,i|
    puts "##{i}: #{entry[0]}"
    entry[1].each { |h| puts "  #{h}" }
  end
end

def beaker_run_type
  type = ENV['TYPE'] || :aio
  type = type.to_sym
end

def sha
  ENV['SHA']
end

def config
  ENV['BEAKER_HOSTS']
end


def get_test_sample
  # This set represents a reasonable sample of puppet acceptance tests,
  # covering a wide range of features and code susceptible to regressions.
  tests = [ 'tests/direct_puppet/cached_catalog_remediate_local_drift.rb',
            'tests/resource/file/content_attribute.rb',
            'tests/environment/environment_scenario-default.rb',
            'tests/puppet_apply_basics.rb',
            'tests/modules/install/basic_install.rb',
            'tests/face/loadable_from_modules.rb',
            'tests/language/functions_in_puppet_language.rb',
            'tests/node/check_woy_cache_works.rb',
            'tests/parser_functions/calling_all_functions.rb',
            'tests/ticket_4622_filebucket_diff_test.rb',
            'tests/pluginsync/4420_pluginfacts_should_be_resolvable_on_agent.rb',
            'tests/ssl/puppet_cert_generate_and_autosign.rb',
            'tests/resource/service/puppet_mcollective_service_management.rb'
          ]

  # Add any tests modified within the last two weeks to the list
  modified = `git log --name-only --pretty="format:" --since 2.weeks ./tests`
  tests += modified.split("\n").reject { |s| s.empty? }.collect { |s| s.sub('acceptance/', '')}

  tests.uniq
end

namespace :ci do
  task :check_env do
    raise(USAGE) unless sha
  end

  namespace :test do
    USAGE = <<-EOS
Requires commit SHA to be put under test as environment variable: SHA='<sha>'.
Also must set BEAKER_HOSTS=config/nodes/foo.yaml or include it in an options.rb for Beaker,
or specify TEST_TARGET in a form beaker-hostgenerator accepts, e.g. ubuntu1504-64a.
You may override the default master test target by specifying MASTER_TEST_TARGET.
You may set TESTS=path/to/test,and/more/tests.
You may set additional Beaker OPTIONS='--more --options'
If testing from git checkouts, you may optionally set the github fork to checkout from using PUPPET_FORK='some-other-puppet-fork' (you may change the HIERA_FORK and FACTER_FORK as well if you wish).
You may also optionally set the git server to checkout repos from using GIT_SERVER='some.git.mirror'.
Or you may set PUPPET_GIT_SERVER='my.host.with.git.daemon', specifically, if you have set up a `git daemon` to pull local commits from.  (You will need to allow the git daemon to serve the repo (see `git help daemon` and the docs/acceptance_tests.md for more details)).
If there is a Beaker options hash in a ./local_options.rb, it will be included.  Commandline options set through the above environment variables will override settings in this file.
EOS

    desc <<-EOS
Run a limited but representative subset of acceptance tests through Beaker and install packages as part of the AIO puppet-agent installation. This task is intended to reduce testing time on a per-commit basis.

  $ env SHA=<full sha> bundle exec rake ci:test:quick
EOS

  task :quick => 'ci:check_env' do
    ENV['TESTS'] = get_test_sample.join(",")
    beaker_test(:aio)
  end

    desc <<-EOS
Run the acceptance tests through Beaker and install packages as part of the AIO puppet-agent installation.
#{USAGE}
EOS
    task :aio => 'ci:check_env' do
      beaker_test(:aio)
    end

    desc <<-EOS
Install puppet as a gem on a predefined set of hosts using Beaker, and run a basic smoke test.

  $ env SHA=<full sha> bundle exec rake:ci:gem
EOS
    task :gem => 'ci:check_env' do
      ENV['TESTS'] = 'setup/gem/pre-suite/010_GemInstall.rb'
      ENV['BEAKER_HOSTS'] = 'config/nodes/gem.yaml'
      beaker_test(:gem)
    end

    desc <<-EOS
Run the acceptance tests through Beaker and install from git on the configuration targets.
#{USAGE}
EOS
    task :git => 'ci:check_env' do
      beaker_test(:git)
    end
  end

  desc "Capture the master and agent hostname from the latest log and construct a preserved_config.yaml for re-running against preserved hosts without provisioning."
  task :extract_preserved_config do
    generate_config_for_latest_hosts
  end

  desc <<-EOS
Run an acceptance test for a given node configuration and preserve the hosts.
Defaults to an aio run, but you can set it to 'git' with TYPE='git'.
#{USAGE}
  EOS
  task :test_and_preserve_hosts => 'ci:check_env'  do
    beaker_test(beaker_run_type, :preserve_hosts => 'always', :__preserve_config__ => true)
  end

  desc "List acceptance runs from the past day which had hosts preserved."
  task :list_preserved do
    preserved = list_preserved_configurations
    print_preserved(preserved)
  end

  desc <<-EOS
Shutdown and destroy any hosts that we have preserved for testing.  These should be reaped daily by scripts, but this will free up resources immediately.
Specify a list of comma separated HOST_NAMES if you have a set of dynamic vcloud host names you want to purge outside of what can be grepped from the logs.
You can go back through the last SECS_AGO logs.  Default is one day ago in secs.
  EOS
  task :release_hosts do
    host_names = ENV['HOST_NAMES'].split(',') if ENV['HOST_NAMES']
    secs_ago = ENV['SECS_AGO']
    release_hosts(host_names, secs_ago)
  end

  task :destroy_preserved_hosts => 'ci:release_hosts' do
    puts "Note: we are now releasing hosts back to the vcloud pooling api rather than destroying them directly.  The rake task for this is ci:release_hosts"
  end

  desc <<-EOS
Rerun an acceptance test using the last captured preserved_config.yaml to skip provisioning.
Or specify a CONFIG_NUMBER from `rake ci:list_preserved`.
Defaults to an aio run, but you can set it to 'git' with TYPE='git'.
  EOS
  task :test_against_preserved_hosts do
    config_number = (ENV['CONFIG_NUMBER'] || 0).to_i
    preserved = list_preserved_configurations
    print_preserved(preserved)
    config_path = preserved[config_number][0]

    puts "Using ##{config_number}: #{config_path}"

    options = {
      :hosts_file => "#{config_path}/preserved_config.yaml",
      :no_provision => true,
      :preserve_hosts => 'always',
    }
    options.merge!(:__delete_options__ => [:pre_suite])

    beaker_test(beaker_run_type, options)
  end
end

task :default do
  sh('rake -T')
end

task :spec do
  sh('rspec lib')
end
