#!/usr/bin/env perl
#ABSTRACT: Run a job in the cluster using NBI::Slurm
#PODNAME: runjob

use v5.12;
use warnings;
use Getopt::Long;
use FindBin qw($RealBin);
use Data::Dumper;
$Data::Dumper::Sortkeys = 1;

if (-e "$RealBin/../dist.ini") {
    say STDERR "[dev mode] Using local lib";
    use lib "$RealBin/../lib";
} 

use NBI::Slurm;
use Cwd;

my $user_home_dir = $ENV{HOME};
my @afterok;
my $user_current_dir = getcwd();
my $username = $ENV{USER};
my $config = load_config("$user_home_dir/.nbislurm.config");
my $version = $NBI::Slurm::VERSION;
my $queue = $config->{'queue'} // 'qib-short,nbi-short';
my $threads = 1;
my $memory = 8000;
my $time = 2;
my $tmpdir = "/tmp";
my $name;
my $email_address = $config->{'email'} // undef;
my $mail_type = $config->{'email_type'} // "none";
my $command;
my $verbose;
my $run;
GetOptions(
    'm|memory=s' => \$memory,
    't|threads=i' => \$threads,
    'q|queue=s' => \$queue,
    'time=s' => \$time,
    'T|tmpdir=s' => \$tmpdir,
    'n|name=s' => \$name,
    'r|run' => \$run,
    'a|email-address=s' => \$email_address,
    'e|mail-type=s' => \$mail_type, # 'BEGIN,END,FAIL,REQUEUE,ALL'
    'after=s' => \@afterok,
    'verbose' => \$verbose,
    'help' => sub { usage() },
);

# Update queue based on time
$queue = update_queue($queue, $time);
$command = join(" ", @ARGV);
usage() unless ( $command);

$name = autoname($command) unless (defined $name);

my $afterok_string = defined @afterok[0] ? "-d afterok:" . join(":", @afterok) : undef;

my $opts = NBI::Opts->new(
    -queue => $queue,
    -threads => $threads,
    -memory => $memory,
    -time   => $time,
    -tmpdir => $tmpdir,
    -email_address => $email_address,
    -email_type => $mail_type,
    -opts => [$afterok_string],
);

if ($verbose) {
    say STDERR "CONFIG:\n", Dumper($config);
    say STDERR $opts->view();
}

my $job = NBI::Job->new(
    -name => $name,
    -command =>  "cd \"$user_current_dir\"",
    -opts => $opts
);
$job->append_command($command);
if ($verbose) {
    say STDERR $job->script();
}

if ($run) {
    if (my $j = $job->run()) {
        say "Job submitted: $j";
    } else {
        say "Job not submitted: $j";
    }
} else {
    say $job->script();
}

sub usage {
    say STDERR <<END;
 ----------------------------------------------------------
  NBI Slurm queuer $version
 ----------------------------------------------------------
 runjob [options] "Command to run"

 Options:
    -n, --name       Job name [optional]
    -q, --queue      Queue name [default: nbi-short]
    -m, --memory     Memory to use [default: 8Gb]
    -t, --threads    Number of threads [default: 1]
    --time           Time string [default: 0d 8h]
    -T, --tmpdir     Temporary directory [default: /tmp]
    -r, --run        Run the job (otherwise, just print the script)
    --verbose        Verbose output
    --help           This help message
 ----------------------------------------------------------
END
    exit;
}

sub autoname {
    my $string = shift;
    my @parts = split(/\s+/, $string);
    my @ints = ["bash", "perl", "python", "python3", "R", "Rscript", "sh", "zsh"];
    # If the first part is in ints, use the second item as name
    if (grep {$_ eq $parts[0]} @ints) {
        return $parts[1];
    } else {
        return $parts[0];
    }
}

sub load_config {
    my $filename = shift;
    my $config = {};
    if (! -e "$filename") {
        say STDERR "Config file not found: $filename";
        return $config;
    }
    open(my $fh, "<", $filename) or die "Cannot open $filename: $!";
    while (<$fh>) {
        chomp;
        next if (/^\s*$/);
        next if (/^#/);
        next if (/^;/);
        my ($key, $value) = split(/=/, $_);
        # discard keys with spaces
        next if ($key =~ /\s/);
        $config->{$key} = $value;
    }
    close  $fh;
    return $config;
}

sub update_queue {
    my ($queue, $time) = @_;
    # If the queue has a star, it can be updated
    if ($queue !~ /\*/) {
        return $queue;
    }
    # Get time in hours
    my $mock_opt = NBI::Opts->new(-time => $time);
    my $time_h = $mock_opt->hours;
    
    if ($time <= 2 ) {
        # Replace * with short
        $queue =~ s/\*/short/g;
    } elsif ($time_h <= 24) {
        # Replace * with medium
        $queue =~ s/\*/medium/g;
    } else {
        # Replace * with long
        $queue =~ s/\*/long/g;
    }
    return $queue;
}

__END__

=pod

=encoding UTF-8

=head1 NAME

runjob - Run a job in the cluster using NBI::Slurm

=head1 VERSION

version 0.4.0

=head1 SYNOPSIS

runjob [options] "Command to run"

=head1 DESCRIPTION

The C<runjob> script allows you to submit a job to the cluster using the NBI::Slurm module. 
It provides a command-line interface for setting the job parameters, including the queue, memory, threads, and execution time.

=head1 OPTIONS

=over 4

=item B<-n, --name>

Specifies the name of the job (optional). If not provided, an automatic name will be generated based on the command being run.

=item B<-q, --queue>

Specifies the queue name for the job. The default value is "nbi-short".
Note that if you put a "*", it will be replaced by 'short', 'medium', or 'long' depending on the time specified with '--time'.

=item B<-m, --memory>

Specifies the amount of memory to use for the job. The default value is 8Gb.

=item B<-t, --threads>

Specifies the number of threads to use for the job. The default value is 1.

=item B<--time>

Specifies the time string for the job. The default value is "0d 8h". The format should be in the format of "Xd Xh Xm" where X represents the number of days (d), hours (h), and minutes (m) respectively.

=item B<-T, --tmpdir>

Specifies the temporary directory for the job. The default value is "/tmp".

=item B<-r, --run>

Runs the job immediately after submitting. If not specified, the script will only print the job script without running it.

=item B<--verbose>

Enables verbose output, displaying additional information about the job and its options.

=item B<--help>

Displays the help message for the script.

=back

=head1 CONFIGURATION

The script will look for a configuration file in the user's home directory at C<~/.nbislurm.config>.

  queue=qib-*,nbi-*
  email=my@address

=head1 EXAMPLES

Submitting a job to the default queue with 4Gb memory and running the job:

    runjob -m 4Gb -r "ls -l"

Submitting a job with a custom name, 2 threads, and running a Python script:

    runjob -n "my-job" -t 2 -r "python script.py"

Printing the job script without running it:

    runjob "echo 'Hello, world!'"

=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
