#!/usr/bin/env ruby

require 'rbbt-util'
require 'rbbt/util/simpleopt'
require 'rbbt/hpc'

#$0 = "rbbt #{$previous_commands*""} #{ File.basename(__FILE__) }" if $previous_commands

options = SOPT.setup <<EOF

Clean error or aborted jobs

$ rbbt slurm clean [options] 

-h--help Print this help
-d--done Done jobs only
-e--error Error jobs only
-a--aborted SLURM aboted jobs
-q--queued Queued jobs only
-j--job* Job ids
-s--search* Regular expression
-t--tail* Show the last lines of the STDERR
-BP--batch_parameters show batch parameters
-dr--dry_run Do not erase anything
EOF

if options[:help]
 if defined? rbbt_usage
  rbbt_usage 
 else
  puts SOPT.doc
 end
 exit 0
end

batch_system = options.delete :batch_system
batch_system ||= 'auto'

HPC::BATCH_MODULE = case batch_system.to_s.downcase
               when 'slurm'
                 HPC::SLURM
               when 'lsf'
                 HPC::LSF
               when 'auto'
                 case $previous_commands.last
                 when 'slurm'
                   HPC::SLURM
                 when 'lsf'
                   HPC::LSF
                 else
                   case Rbbt::Config.get(:batch_system, :batch, :batch_system, :hpc, :HPC, :BATCH).to_s.downcase
                   when 'slurm'
                     HPC::SLURM
                   when 'lsf'
                     HPC::LSF
                   else
                     case ENV["BATCH_SYSTEM"].to_s.downcase
                     when 'slurm'
                       HPC::SLURM
                     when 'lsf'
                       HPC::LSF
                     end
                   end
                 end
               end

raise ParameterException.new("Could not detect batch_system: #{Misc.fingerprint batch_system}") if HPC::BATCH_MODULE.nil?

Log.severity = 4
done, error, aborted, queued, jobid, search, tail, batch_parameters, dry_run = options.values_at :done, :error, :aborted, :queued, :job, :search, :tail, :batch_parameters, :dry_run

workdir = File.expand_path('~/rbbt-batch')
Path.setup(workdir)

running_jobs = begin
                squeue_txt = HPC::BATCH_MODULE.job_status
                squeue_txt.split("\n").collect{|l| l.to_i.to_s}
               rescue
                Log.warn "Cannot determine if jobs are running, they will seem to be all alive (Job ID in green)"
                squeue_txt = nil
                $norunningjobs = true
                []
               end

if squeue_txt
 job_nodes = {}
 squeue_txt.split("\n").each do |line|
  parts = line.strip.split(/\s+/)
  job_nodes[parts.first] = parts.last.split(",")
 end
else
 job_nodes = nil
end

count = 0
workdir.glob("**/command.batch").sort_by{|f| File.mtime(f)}.each do |fcmd|
 dir = File.dirname(fcmd)
 command_txt = Open.read(fcmd)

 if m = command_txt.match(/#CMD: (.*)/)
  cmd = m[1]
 else
  cmd = nil
 end

 if m = command_txt.match(/# Run command\n(.*?)\n/im)
  exe = m[1]
 else
  exe = nil
 end

 if m = command_txt.match(/^CONTAINER_DIR=(.*)/)
  container_home = m[1]
 else
  container_home = nil
 end

 if m = command_txt.match(/^BATCH_SYSTEM=(.*)/)
  job_batch_system = m[1].downcase
 else
  job_batch_system = nil
 end

 different_system = job_batch_system != batch_system

 if File.exists?(fid = File.join(dir, 'job.id'))
  id = Open.read(fid).chomp
 else
  id = nil
 end

 if File.exists?(fstatus = File.join(dir, 'exit.status'))
  exit_status = Open.read(fstatus).to_i
 else
  exit_status = nil
 end

 if File.exists?(fstatus = File.join(dir, 'job.status'))
   fstatus_txt = Open.read(fstatus) 
   begin
     if job_batch_system == "lsf"
       nodes = Open.read(fstatus).split("\n").last.split(/\s+/)[5].split(",")
     else
       nodes = Open.read(fstatus).split("\n").last.split(/\s+/).last.split(",")
     end
   rescue
     nodes = []
   end
 elsif job_nodes[id]
  nodes = job_nodes[id]
 else
  nodes = []
 end

 if File.exists?(File.join(dir, 'std.out'))
  outt = File.mtime File.join(dir, 'std.out')
  errt = File.mtime File.join(dir, 'std.err')
  time_diff = Time.now - [outt, errt].max
 end

 fdep = File.join(dir, 'dependencies.list')
 deps = Open.read(fdep).split("\n") if File.exists?(fdep)

 fcadep = File.join(dir, 'canfail_dependencies.list')
 cadeps = Open.read(fcadep).split("\n") if File.exists?(fcadep)

 aborted = error = true if aborted.nil? && error.nil?
 #if done || error || aborted || running || queued || jobid || search
 #  select = false
 #  select = true if done && exit_status && exit_status.to_i == 0 
 #  select = true if error && exit_status && exit_status.to_i != 0 
 #  select = true if aborted && (exit_status.nil? && ! running_jobs.include?(id))
 #  select = select && jobid.split(",").include?(id) if jobid 
 #  select = select && cmd.match(/#{search}/) if search
 #  next unless select 
 #end

 if done || error || aborted || queued || jobid 
   select = false
   select = true if done && exit_status == 0
   select = true if error && exit_status && exit_status != 0 
   select = true if aborted && (exit_status.nil? && ! running_jobs.include?(id))
   is_running =  exit_status.nil? && ( (running_jobs.include?(id) && (!deps || (running_jobs & deps).empty?)) || different_system )
   select = true if queued && deps && (running_jobs & deps).any? || queued && is_running && nodes.empty?
   select = true if jobid && jobid.split(",").include?(id)
   select = select && cmd.match(/#{search}/) if search
   next unless select 
 elsif search
   select = false
   select = true if search && cmd.match(/#{search}/)
   next unless select 
 end


 puts Log.color(:yellow, "**ERASING**")
 puts Log.color :blue, dir
 puts Log.color(:magenta, "Creation: ") << File.mtime(File.join(dir, 'command.batch')).to_s
 puts Log.color(:magenta, "Done: ") << File.mtime(File.join(dir, 'exit.status')).to_s if File.exist?(File.join(dir, 'exit.status'))
 puts Log.color(:magenta, "Exec: ") << (exe || "Missing")
 puts Log.color(:magenta, "CMD: ") << (Log.color(:yellow, cmd) || "Missing")
 puts Log.color(:magenta, "HOME: ") << Log.color(:yellow, container_home) if container_home
 puts Log.color(:magenta, "Job ID: ") << (exit_status ? (exit_status == 0 ? Log.color(:green, "Done") : Log.color(:red, "Error")) + " (#{ id })" :  (running_jobs.include?(id) || $norunningjobs ? Log.color(:green, id) : Log.color(:red, id) ))
 puts Log.color(:magenta, "Dependencies: ") << deps * ", " if deps
 puts Log.color(:magenta, "Dependencies (can fail): ") << cadeps * ", " if cadeps
 puts Log.color(:magenta, "Nodes: ") << nodes * ", "
 puts Log.color(:magenta, "Output: ") << File.exists?(File.join(dir, 'std.out')).to_s << (id.nil? ? "" : " (last update " + Misc.format_seconds(time_diff) + " ago)")

 if options[:batch_parameters]
   puts Log.color(:magenta, "BATCH parameters: ")
   case job_batch_system 
   when 'slurm'
     puts Log.color :blue, CMD.cmd('grep "^#SBATCH" |tail -n +6', :in => Open.read(fcmd)).read.strip
   when 'lsf'
     puts Log.color :blue, CMD.cmd('grep "^#BSUB" |tail -n +6', :in => Open.read(fcmd)).read.strip
   end
 end

 if tail && File.exists?(File.join(dir, 'std.err'))
   if exit_status && exit_status != 0
     puts Log.color(:magenta, "First error or exception found: ")
     puts CMD.cmd("grep -i -w 'error\\|[a-z]*exception' #{File.join(dir, 'std.err')} -A #{tail.to_i} |head -n #{tail.to_i}", :no_fail => true).read
   elsif exit_status
     puts Log.color(:magenta, "Completed jobs: ")
     puts CMD.cmd("grep -i -w 'Completed step' #{File.join(dir, 'std.err')} | grep -v 'Retrying dep.' | tail -n #{tail.to_i}", :no_fail => true).read
   else
     puts Log.color(:magenta, "Log tail: ")
     puts CMD.cmd("tail -n #{tail.to_i} #{File.join(dir, 'std.err')}").read
   end
 end

 count += 1

 Open.rm_rf dir unless dry_run
end

puts 
puts "Found #{count} jobs"