#! /usr/bin/env ruby

# frozen_string_literal: true

LKP_SRC = ENV['LKP_SRC'] || File.dirname(File.dirname(File.realpath($PROGRAM_NAME)))

require "#{LKP_SRC}/lib/local_pack"
require "#{LKP_SRC}/lib/job2sh"
require "#{LKP_SRC}/lib/scheduler_client"
require 'optparse'
require 'yaml'

opt_set_key_value = {}
opt_cmdline_defaults = {}
opt_output_dir = nil
opt_auto_define_files = false
opt_include_yamls = []
opt_monitor = false
opt_monitor_query = {}
opt_my_queue = false
actions = ['output', 'stop']
result_roots = []
nr_run = 1
do_pack = true

options = OptionParser.new do |opts|
  opts.banner = 'Usage: submit [options] job1.yaml job2.yaml ...'

  opts.separator '       submit test jobs to the scheduler'
  opts.separator ''
  opts.separator 'options:'

  opts.on("-s 'KEY: VALUE'", "--set 'KEY: VALUE'", 'add YAML hash to job') do |key_value|
    k, v = key_value.sub(' ', '').split(':', 2)
    opt_set_key_value[k] = v
  end

  opts.on('-o DIR', '--output DIR', 'save job yaml to DIR/') do |dir|
    if File.file? dir
      puts "Please input directory for job save yaml after '-o'"
      exit 1
    end
    opt_output_dir = dir
    Dir.mkdir(dir) unless File.directory? dir
  end

  opts.on('-a', '--auto-define-files', 'auto add define_files') do
    opt_auto_define_files = true
  end

  opts.on('--no-pack', "don't do pack, just use the last one") do
    do_pack = false
  end

  opts.on('-i include.yaml', '--include include.yaml', 'include other job yamls') do |yaml|
    opt_include_yamls << yaml
  end

  opts.on('-c', '--connect', 'auto connect to the host') do
    actions << 'connect'
  end

  opts.on('-r', '--result', 'mirror job result dir') do
    actions << 'mirror_result'
  end

  opts.on('-n job_number', '--number job_number', 'number to submit job') do |number|
    nr_run = number.to_i
    nr_run = 1 if nr_run < 1
  end

  opts.on('-m', '--monitor', "monitor job status: use -m 'KEY: VALUE' to add rule") do
    opt_monitor = true
    k, v = ARGV[0].sub(' ', '').split(':', 2) if ARGV[0]
    if (k && !k.empty?) && (v && !v.empty?)
      opt_monitor_query[k] = v
      ARGV.shift
    end
  end

  opts.on('--my-queue', 'add to my queue') do
    opt_my_queue = true
  end

end

ARGV_CLONE = ARGV.clone
options.parse!(ARGV)

seen_job_yaml = false
ARGV.delete_if do |arg|
  if arg.index '='
    k, v = arg.split('=', 2)
    if seen_job_yaml
      opt_set_key_value[k] = v
    else
      opt_cmdline_defaults[k] = v
    end
    true
  else
    seen_job_yaml = true if arg =~ /\.yaml$/
    false
  end
end

if ARGV.size.zero?
  puts(options)
  exit
end

job_ids = []
job_hash_list = []

def prioritize_include_yaml(include_yamls, jobfile)
  default_yamls, override_yamls = [], []
  jobfile_index = ARGV_CLONE.index(jobfile)
  include_yamls.each do |yaml|
    if ARGV_CLONE.index(yaml) < jobfile_index
      default_yamls << yaml
    else
      override_yamls << yaml
    end
  end
  return default_yamls, override_yamls
end

def find_jobfile(jobfile)
  return jobfile if File.exist?(jobfile)

  search_paths = [Dir.pwd, File.join(LKP_SRC, 'jobs')]
  search_paths.each do |search_path|
    search_jobfile = File.join(search_path, jobfile)
    return search_jobfile if File.exist?(search_jobfile)
  end
  puts("Cannot find job #{jobfile} in directories:"); puts search_paths
  exit 1
end

def find_jobfiles(jobfile_list)
  search_jobfile_list = []
  jobfile_list.each do |jobfile|
    search_jobfile_list << find_jobfile(jobfile)
  end
  return search_jobfile_list
end

submit_id = %x(uuidgen).chomp
puts "submit_id=#{submit_id}"

ARGV.each do |jobfile|
  default_yamls, override_yamls = prioritize_include_yaml(opt_include_yamls, jobfile)
  jobfile = find_jobfile(jobfile)
  jobs = Job2sh.new
  jobs.cmdline_defaults = opt_cmdline_defaults
  jobs.overrides = opt_set_key_value
  jobs.default_yamls = find_jobfiles(default_yamls)
  jobs.override_yamls = find_jobfiles(override_yamls)
  jobs.load(jobfile, true) || next
  jobs[:expand_params] = true
  jobs['testbox'] = opt_set_key_value['testbox'] if opt_set_key_value['testbox']
  jobs['tbox_group'] = tbox_group(jobs['testbox']) if jobs.include?('testbox')
  jobs['node_roles'] ||= 'server client' if jobs['cluster']
  jobs['submit_id'] = submit_id
  jobs.each_jobs do |job|
    raise 'Please configure SCHED_HOST' unless job['SCHED_HOST']
    raise 'Please configure SCHED_PORT' unless job['SCHED_PORT']

    job['queue'] = "#{jobs['tbox_group']}~#{ENV['USER']}" if opt_my_queue

    job.add_pp
    job.add_monitors
    job.add_define_files if opt_auto_define_files

    # get job shell function
    sh_run_job = job.sh_run_job
    sh_extract_stats = job.sh_extract_stats
    sh_define_files = job.sh_define_files
    sh_on_fail = job.sh_on_state(state='on_fail')

    sh_hash = {
      'job2sh' => {
        'run_job' => sh_run_job,
        'extract_stats' => sh_extract_stats,
        'define_files' => sh_define_files
      }
    }

    if sh_on_fail
      sh_hash['job2sh']['on_fail'] = sh_on_fail
    end

    # merge job info
    job_hash = job.to_hash
    job_hash_list << job_hash
    job_hash = job_hash.merge(sh_hash)

    # save job to yaml
    if opt_output_dir
      prefix = File.join(opt_output_dir, File.basename(jobfile, '.yaml'))
      unit_jobfile = prefix + '-' + job.path_params[0..180] + '.yaml'
      job.save unit_jobfile
      puts "#{jobfile} => #{unit_jobfile}"
      next
    end

    scheduler_client = SchedulerClient.new(job['SCHED_HOST'], job['SCHED_PORT'])

    # submit job
    job_hash['nr_run'] = nr_run

    # do local pack
    # modified files since the last tag
    # untracked files
    pkg_data = {}

    unless job_hash.key?('pkg_data')
      lkp_repos = [ENV['LKP_SRC']]
      lkp_repos.insert(-1, ENV['LKP_SRC2']) if ENV['LKP_SRC2']

      lkp_repos.each do |repo|
        repo_name = File.basename(repo)
        do_package = PackChange.new(repo, do_pack)
        tag, md5, content = do_package.pack_source

        pkg_data[repo_name] = {
          'tag' => tag,
          'md5' => md5,
          'content' => content
        }
      end
    end

    # init scheduler client
    # add pkg_data to job_hash
    job_hash['pkg_data'] = pkg_data unless pkg_data.empty?
    job_json = job_hash.to_json
    nr_run.times do
      messages = scheduler_client.submit_job(job_json)
      JSON.parse(messages).each do |msg|
        if msg['message'].empty?
          result_roots << msg['result_root']
          job_ids << msg['job_id'].to_s
          puts("submit #{jobfile}, got job id=#{msg['job_id']}")
        else
          opt_monitor = false
          puts("submit #{jobfile} failed, got job id=#{msg['job_id']}, error: #{msg['message']}")
          if msg['message'].include?('Failed to verify the account.')
            puts 'The submitted account information is as follows:'
            puts "my_name: #{job['my_name']}"
            puts "my_email: #{job['my_email']}"
            puts "my_account: #{job['my_account']}"
            puts "my_token: #{job['my_token']}"
          end
        end
      end
    end
  end
end

if opt_monitor
  job_hash_list[0].delete('define_files')
  opt_monitor_query.merge!({'job_id' => job_ids})
  cmd = "#{LKP_SRC}/sbin/monitor -f '#{opt_monitor_query.to_json}' -j '#{job_hash_list[0].to_json}' -a #{actions.join(',')} -r #{result_roots.join(',')} -s 'job_state: extract_finished'"
  exec cmd
end
