#!/usr/bin/env perl
#ABSTRACT: Wait for selected jobs to exit the queue
#PODNAME: waitjobs

use v5.12;
use warnings;
use Getopt::Long;
use FindBin qw($RealBin);
use Data::Dumper;
use Term::ANSIColor qw(:constants);
if (-e "$RealBin/../dist.ini") {
    say STDERR "[dev mode] Using local lib"  if ($ENV{"DEBUG"});
    use lib "$RealBin/../lib";
} 
use NBI::Slurm;
use NBI::Queue;
$Data::Dumper::Sortkeys = 1;

my $opt_user = $ENV{USER};
my $opt_name = '.+';
my $opt_status = '';
my $opt_verbose_bool = 0;
my $opt_debug_bool = 0;
my $refresh_time = 20;
# Check if this script is running as a job, in case save the job id
my $self_job_id = $ENV{SLURM_JOB_ID} // undef;
GetOptions(
    'u|user=s' => \$opt_user,
    'n|name=s' => \$opt_name,
    'r|refresh=i' => $refresh_time,
    'verbose' => \$opt_verbose_bool,
    'version' => sub { say "waitjobs v", $NBI::Slurm::VERSION; exit },
    'debug' => \$opt_debug_bool,
    'help' => sub { usage() },
);

my $queue = NBI::Queue->new(
    -username => $opt_user,
    -name => $opt_name,
);
# If $self_job_id is anq integer, remove it from the queue
if (defined $self_job_id and $self_job_id =~ /^\d+$/) {
    print STDERR "Ignoring current job id $self_job_id\n";
    $queue->remove($self_job_id);
}

my @starting_ids = @{$queue->ids()};
print STDERR "[waitjob] There are ", 
  GREEN, BOLD, $queue->len(), RESET  " jobs for ",
  GREEN, BOLD, $opt_user,     RESET, " with name matching ",
  GREEN, BOLD, $opt_name,     RESET, "\n";

while (scalar @starting_ids > 0) {
    my $current_jobs = NBI::Queue->new();
    my @current_ids = @{$current_jobs->ids()};

    # Remove from @starting_ids all the jobs that are NOT in @current_ids
    my @new_starting_ids = ();
    for my $starting_id (@starting_ids) {
        if (grep {$_ eq $starting_id} @current_ids) {
            push(@new_starting_ids, $starting_id);
        }
    }
    @starting_ids = @new_starting_ids;
    if ($opt_verbose_bool) {
        say STDERR "Waiting for ", scalar @starting_ids, " jobs to finish";
    }
    sleep $refresh_time;
}


sub usage {
    say STDERR<<END;

  Usage: waitjobs [options] [job_name]
  ------------------------------------
    Wait for selected jobs to exit the queue

  Options:
   -u, --user       default current user
   -n, --name       default '.+' (matches any name)
   -r, --refresh    check every 20 seconds
    
END
    exit;
}

__END__

=pod

=encoding UTF-8

=head1 NAME

waitjobs - Wait for selected jobs to exit the queue

=head1 VERSION

version 0.8.1

=head1 SYNOPSIS

waitjobs [options]

=head1 DESCRIPTION

This script waits for selected jobs to exit the queue. It continuously checks the status of the jobs and waits until all the specified jobs are completed.

=head1 OPTIONS

=over 4

=item B<-u, --user <username>>

Specify the username to wait for jobs from. By default, it uses the current user.

=item B<-n, --name <pattern>>

Specify the name pattern to filter the jobs. Only jobs with matching names will be considered. The default pattern is '.+' (matches any name).

=item B<-r, --refresh <time>>

Specify the refresh time in seconds for checking the job status. The script will wait for this duration before checking the job status again. The default refresh time is 20 seconds.

=item B<--verbose>

Display verbose output, including the number of jobs being waited for.

=item B<--debug>

Enable debug mode, printing the output of the squeue command.

=item B<--help>

Print the help message and exit.

=back

=head1 EXAMPLES

=over 4

=item B<Example 1:>

Wait for all your jobs with a specific name to exit the queue:

  waitjobs -n myjob

=item B<Example 2:>

Wait for jobs from a specific user to exit the queue:

  waitjobs -u username

=item B<Example 3:>

Wait for jobs with a specific name from a specific user:

  waitjobs -u username -n myjob

=item B<Example 4:>

Wait for jobs with a specific status, refreshing every 10 seconds:

  waitjobs -n myjob -r 10 -u username -s RUNNING

=back

=head1 AUTHOR

Andrea Telatin <proch@cpan.org>

=head1 COPYRIGHT AND LICENSE

This software is Copyright (c) 2023 by Andrea Telatin.

This is free software, licensed under:

  The MIT (X11) License

=cut
