#!/usr/bin/env ruby

$LOAD_PATH.unshift File.expand_path('../lib', __dir__)
$stdout.sync = true

require 'chronic_duration'
require 'fileutils'
require 'influxdb'
require 'optimist'
require 'pathname'
require 'rainbow'
require 'run_k6'
require 'test_info'
require 'time'
require 'tmpdir'
require 'result_writer'

k6_dir = Pathname.new(File.expand_path('../k6', __dir__)).relative_path_from(Dir.pwd)

gpt_version = '2.15.1'
puts Rainbow("GitLab Performance Tool (GPT) v#{gpt_version} - Performance test runner for GitLab environments based on k6").color(230, 83, 40)

opts = Optimist.options do
  version gpt_version
  banner "\nDocumentation: https://gitlab.com/gitlab-org/quality/performance/blob/main/docs/README.md"
  banner "\nUsage: run-k6 [options]"
  banner "Options:"
  opt :environment, "Name of Environment Config file in environments directory that the test(s) will be run with. Alternative filepath can also be given.", type: :string, required: true
  opt :options, "Name of Options Config file in options directory that the test(s) will be run with. Alternative filepath can also be given.", type: :string, default: "20s_2rps.json"
  opt :tests, "Names of Test files or directories to run with. When directory given tests will be recursively added from api, web and git subdirs.", type: :strings, default: ["tests"]
  opt :scenarios, "Include any tests inside the test directory's scenarios subfolder when true.", type: :flag, default: false
  opt :quarantined, "Include any tests inside the test directory's quarantined subfolder when true.", type: :flag, default: false
  opt :vulnerabilities, "Include any tests that test vulnerability report/api", type: :flag, default: false
  opt :excludes, "List of words used to exclude tests by matching against their names.", type: :strings
  opt :unsafe, "Include any tests that perform unsafe requests (POST, PUT, DELETE, PATCH)", type: :flag, default: false
  opt :rate_limits, "Allow GPT to disable certain rate limits before the tests and restore limits after the test run.", type: :flag, default: false
  opt :experimental, "Include any tests inside the test directory's experimental subfolder when true.", type: :flag, default: false
  opt :unattended, "Skip all user prompts and run through tests automatically.", type: :flag, default: false
  opt :influxdb_url, "URL of an Influx DB server where GPT can optionally upload test run statistics.", type: :string, default: ENV['GPT_INFLUXDB_URL'] || nil
  opt :help, 'Show this help message'
  opt :version, 'Print version and exit'
  banner "\nEnvironment Variable(s):"
  banner "  ACCESS_TOKEN                A valid GitLab Personal Access Token for the specified environment. The token should come from a User that has admin access for the project(s) to be tested and have all permissions set. (Default: nil)"
  banner "  GPT_DEBUG                   Shows debug output when set to true. (Default: nil)"
  banner "  GPT_SKIP_RETRY              Skip failed test retry when set to true. (Default: nil)"
  banner "  GPT_TTFB_P95                Add TTFB 95 to the test results output. (Default: nil)"
  banner "  GPT_RESULTS_SORT_BY_FAILED  Sort GPT results by failed tests first. (Default: nil)"
  banner "\nExamples:"
  banner "  Run all Tests with the 60s_200rps Options file against the 10k Environment:"
  banner "    #{ENV['GPT_DOCKER'] ? 'docker run -it gitlab/gitlab-performance-tool' : $PROGRAM_NAME} --environment 10k.json --options 60s_200rps.json"
  banner "  Run all API Tests with the 60s_200rps Options file against the 10k Environment:"
  banner "    #{ENV['GPT_DOCKER'] ? 'docker run -it gitlab/gitlab-performance-tool' : $PROGRAM_NAME} --environment 10k.json --options 60s_200rps.json --tests api"
  banner "  Run a specific Test with the 60s_200rps Options file against the 10k Environment:"
  banner "    #{ENV['GPT_DOCKER'] ? 'docker run -it gitlab/gitlab-performance-tool' : $PROGRAM_NAME} --environment 10k.json --options 60s_200rps.json --tests api_v4_groups_projects.js"
end

k6_path = RunK6.setup_k6

# Variables
env_file = Dir.glob([opts[:environment], "#{ENV['GPT_DOCKER_CONFIG_DIR'] || ''}/environments/#{opts[:environment]}", "#{k6_dir}/#{opts[:environment]}", "#{k6_dir}/config/environments/#{opts[:environment]}"])[0]
raise "Environment config file '#{opts[:environment]}' not found as given or in default folder. Exiting..." unless File.file?(env_file.to_s)

options_file = Dir.glob([opts[:options], "#{ENV['GPT_DOCKER_CONFIG_DIR'] || ''}/options/#{opts[:options]}", "#{k6_dir}/#{opts[:options]}", "#{k6_dir}/config/options/#{opts[:options]}"])[0]
raise "Options config file '#{opts[:options]}' not found as given or in default folder. Exiting..." unless File.file?(options_file.to_s)

env_vars = RunK6.setup_env_vars(k6_dir:, env_file:, options_file:)
gitlab_env_settings = GPTCommon.get_env_settings(env_url: env_vars['ENVIRONMENT_URL'])

RunK6.check_large_projects_visibility(env_vars:) unless ENV['GPT_SKIP_VISIBILITY_CHECK'] || env_vars['GPT_LARGE_PROJECT_CHECK_SKIP']

opts[:vulnerabilities] = true unless env_vars['ENVIRONMENT_VULNERABILITIES_GROUP'].nil?

if ENV['GPT_START_SLEEP']
  puts "Sleeping for #{ENV['GPT_START_SLEEP']} seconds before starting..."
  sleep ENV['GPT_START_SLEEP'].to_i
end

start_time = Time.now

begin
  tests = RunK6.get_tests(k6_dir:, opts:, env_vars:, gitlab_env_settings: gitlab_env_settings.dup)
  RunK6.prepare_tests(tests:, env_vars:)

  results_home = ENV['GPT_DOCKER_RESULTS_DIR'] || Pathname.new(File.expand_path('../results', __dir__)).relative_path_from(Dir.pwd)
  results_file_prefix = ENV['GPT_RESULTS_PATH_PREFIX'] || "#{env_vars['ENVIRONMENT_NAME']}_#{env_vars['ENVIRONMENT_VERSION'].match?(/\d+\.\d+\.\d+.*/) ? "v#{env_vars['ENVIRONMENT_VERSION'].tr('.', '-')}_" : ''}#{start_time.strftime('%Y-%m-%d_%H%M%S')}"
  results_dir = File.join(results_home, results_file_prefix)
  results_output_file = File.join(results_dir, "#{results_file_prefix}_results_output.log")
  FileUtils.mkdir_p(results_dir)
  puts "\nSaving all test results to #{results_dir}"

  aggregated_results = []
  aggregated_success = true

  # Run tests
  test_redo = false
  tests.each do |test_file|
    status, output = RunK6.run_k6(k6_path:, opts:, env_vars:, options_file:, test_file:, results_dir:, gpt_version:)

    File.open(results_output_file, 'a') do |out_file|
      out_file.puts output
    end

    if !status && !test_redo && ENV['GPT_SKIP_RETRY'] != 'true'
      warn Rainbow("Test failed. Retrying...").yellow
      test_redo = true

      if ENV['GPT_RETRY_SLEEP']
        puts "Sleeping for #{ENV['GPT_RETRY_SLEEP']} seconds before retrying..."
        sleep ENV['GPT_RETRY_SLEEP'].to_i
      end

      redo
    end

    aggregated_results << RunK6.get_test_results(test_file:, status:, output:, test_redo:)
    aggregated_success &&= status

    # Reset redo flag
    test_redo = false

    sleep ENV['GPT_SLEEP_BETWEEN'].to_i if ENV['GPT_SLEEP_BETWEEN'] && !test_file.equal?(tests.last)
  rescue Interrupt
    warn Rainbow("Caught the interrupt. Stopping.").yellow
    exit
  rescue ArgumentError => e
    warn Rainbow(e).yellow
    warn Rainbow("GPT v#{gpt_version} is a major release and contains breaking changes in regards to configuration and test data. If you haven't already, please refer to the latest release notes and documentation.").yellow
    next
  rescue NoMethodError => e
    warn Rainbow("Test failed and output couldn't be parsed: \n#{e}").yellow
    next
  rescue StandardError => e
    warn Rainbow("Test failed: #{e}\n#{e.class}:#{e.backtrace}").red
    aggregated_success = false
  end
  aggregated_score = RunK6.get_results_score(results: aggregated_results, env_vars:)

  if aggregated_results.empty?
    puts "No tests reported any results. Exiting..."
    exit
  end

  # Process test times
  end_time = Time.now
  run_time = (end_time - start_time).round(2)

  puts Rainbow("All k6 tests have finished after #{run_time}s!").green

ensure
  RunK6.restore_environment_rate_limit_settings(env_vars:, gitlab_env_settings:) if opts[:rate_limits]
end

# Output known issues
tests_with_issues = TestInfo.get_tests_info(tests).select { |test_info| !test_info[:issues].nil? }
unless tests_with_issues.empty? || ENV['GPT_SKIP_KNOWN_ISSUES'] == 'true'
  puts "\n█ Known issues\n\nNote that the following endpoints below have known issues. These tests have either been run with a custom lower threshold limit applied or are quarantined until the issue is fixed:\n\n"
  tp.set(:max_width, 150)
  tp(tests_with_issues, :name, :issues)
  puts "\nFull list of issues found both past and present can be found here: https://gitlab.com/gitlab-org/gitlab/-/issues?label_name%5B%5D=Quality%3Aperformance-issues\n"
end

# Output and save test results
# JSON
results_hash = {
  "name" => env_vars['ENVIRONMENT_NAME'],
  "version" => env_vars['ENVIRONMENT_VERSION'],
  "revision" => env_vars['ENVIRONMENT_REVISION'],
  "gpt_version" => gpt_version,
  "option" => File.basename(options_file, '.json'),
  "date" => start_time.strftime('%F'),
  "time" => {
    "start" => start_time.utc.strftime('%T %Z'),
    "start_epoch" => (start_time.to_f * 1000).to_i,
    "end" => end_time.utc.strftime('%T %Z'),
    "end_epoch" => (end_time.to_f * 1000).to_i,
    "run" => run_time
  },
  "overall_result" => aggregated_success,
  "overall_result_score" => aggregated_score,
  "test_results" => aggregated_results
}

writers = [ResultWriter::Json, ResultWriter::Csv, ResultWriter::Txt]

# Main results
result_writers = writers.map { |w| w.new(results_dir, results_file_prefix) }
result_files = [results_output_file] + result_writers.map(&:path)
result_writers.each { |r| r.write(results_hash) }

# Failed results
failed_results_hash = results_hash.dup
failed_results_hash["test_results"] = aggregated_results.select { |test| test['result'] == false }
unless failed_results_hash["test_results"].empty? || ENV['GPT_SKIP_FAILED_RESULT_REPORT']
  failed_results_dir = "#{results_dir}/failed_test_results"
  failed_results_file_prefix = "#{results_file_prefix}_failed"
  FileUtils.mkdir_p(failed_results_dir)
  failed_result_writers = writers.map { |w| w.new(failed_results_dir, failed_results_file_prefix) }
  failed_result_writers.each { |r| r.is_a?(ResultWriter::Txt) ? r.write(failed_results_hash, true) : r.write(failed_results_hash) }
  failure_report = ResultWriter::MD.new(failed_results_dir, "#{results_file_prefix}_failure_evaluation_report")
  failure_report.write(failed_results_hash)
  puts "\n█ Failed Results files\n\n#{failed_result_writers.map(&:path).join("\n")}"
  puts "\n█ Failure Evaluation Report\n\n#{failure_report.path}"
end

puts "\n█ Full Results files\n\n#{result_files.join("\n")}"

if opts[:influxdb_url]
  influxdb_report, message = InfluxDB.write_data(opts[:influxdb_url], results_hash)
  warn Rainbow("\nFailed to upload test run statistics to InfluxDB URL #{opts[:influxdb_url]} - #{message}.").red unless influxdb_report
end

abort("\n" + Rainbow("One or more tests have failed...").red) unless aggregated_success || ENV['GPT_IGNORE_RESULT'] == 'true'
