#!/bin/env perl
###
### pstamp_job_run.pl
###
###     Run a given postage stamp Job
###

use warnings;
use strict;

use Sys::Hostname;
use Getopt::Long qw( GetOptions :config auto_help auto_version gnu_getopt );
use Carp;
use File::Basename;
use File::Copy;
use File::Temp qw(tempfile tempdir);
use Digest::MD5::File qw( file_md5_hex );
use PS::IPP::PStamp::RequestFile qw( :standard );
use PS::IPP::PStamp::Job qw( :standard );
use IPC::Cmd 0.36 qw( can_run run );
use POSIX;

use PS::IPP::Metadata::Config;
#use PS::IPP::Metadata::Stats;
use PS::IPP::Metadata::List qw( parse_md_list );

use PS::IPP::Config qw( :standard );

my ($job_id, $redirect_output, $outputBase, $rownum, $jobType, $options); 
my ($verbose, $dbname, $dbserver, $no_update, $save_temps);

GetOptions(
    'job_id=s'          =>  \$job_id,
    'job_type=s'        =>  \$jobType,
    'rownum=s'          =>  \$rownum,
    'output_base=s'     =>  \$outputBase,
    'options=s'         =>  \$options,
    'redirect-output'   =>  \$redirect_output,
    'dbname=s'          =>  \$dbname,
    'dbserver=s'        =>  \$dbserver,
    'verbose'           =>  \$verbose,
    'no-update'         =>  \$no_update,
    'save-temps'        =>  \$save_temps,
);


my $host = hostname();
if ($verbose) {
    print "\n\n";
    print "Starting script $0 on $host\n\n";
}

die "job_id is required" if !$job_id;

my_die( "job_type is required", $job_id, $PS_EXIT_PROG_ERROR) if !$jobType;
my_die("rownum is required", $job_id, $PS_EXIT_PROG_ERROR) if !$rownum;
my_die("output_base is required", $job_id, $PS_EXIT_PROG_ERROR) if !$outputBase;

$options = 1 if !$options;


my $ipprc = PS::IPP::Config->new(); # IPP Configuration
if ($redirect_output) {
    my $logDest = "$outputBase.log";
    $ipprc->redirect_output($logDest)
        or my_die ("unable to redirect output to $logDest", $job_id, $PS_EXIT_UNKNOWN_ERROR);
}


if (!$dbserver) {
    $dbserver =  metadataLookupStr($ipprc->{_siteConfig}, 'PS_DBSERVER');
}

my $missing_tools;

my $pstamptool = can_run('pstamptool')  or (warn "Can't find pstamptool"  and $missing_tools = 1);
my $ppstamp    = can_run('ppstamp') or (warn "Can't find ppstamp" and $missing_tools = 1);
my $pstamp_get_image_job    = can_run('pstamp_get_image_job.pl') or (warn "Can't find pstamp_get_image_job.pl" and $missing_tools = 1);
my $dqueryparse = can_run('dqueryparse.pl') or (warn "Can't find dqueryparse.pl" and $missing_tools = 1);
my $streaksreplace = can_run('streaksreplace')  or (warn "Can't find streaksreplace"  and $missing_tools = 1);
my $magicdstool = can_run('magicdstool')  or (warn "Can't find magicdstool"  and $missing_tools = 1);

if ($missing_tools) {
    my_die("Can't find required tools", $job_id, $PS_EXIT_CONFIG_ERROR);
}

my $mdcParser = PS::IPP::Metadata::Config->new; # Parser for metadata config files

my $jobStatus;
if ($jobType eq "stamp") {
    my $params = read_params_file($outputBase);

    my $argString;
    $argString = $params->{job_args};
    my $stage = $params->{stage};

    # XXX: should we do any other sanity checking?
    my_die("argument list is empty", $job_id, $PS_EXIT_DATA_ERROR) if !$argString;

    my $nan_masked = 1;
    my $muggle = 0;
    if (!$params->{magicked}) {
        $nan_masked = 0;
    } elsif ($params->{magicked} and ($options & ($PSTAMP_REQUEST_UNCENSORED | $PSTAMP_REQUIRE_UNCENSORED))) {
        # Attempt to find or create a muggle image
        $nan_masked = 0;
        $muggle = 1;
    }

    my $image = $params->{image};
    my $mask;
    my $variance;
    my $fileArgs = " -file $params->{image}";
    my @file_list = ($params->{image});
    
    if ($nan_masked or ($options & $PSTAMP_SELECT_MASK)) {
        $mask = $params->{mask};
        $fileArgs .= " -mask $mask";
        push @file_list, $mask;
    }
    if ($options & $PSTAMP_SELECT_VARIANCE) {
        $variance = $params->{weight};
        $fileArgs .= " -variance $variance";
        push @file_list, $variance;
    }

    if ($params->{astrom}) {
        $argString .= " -astrom $params->{astrom}";
        push @file_list, $params->{astrom};
    }

    # check that actual input files exist
    check_files($PSTAMP_GONE, @file_list);

    # find our output directory
    my $outdir = dirname($outputBase);
    my ($tmpImage, $tmpMask, $tmpVariance, $tmproot);
    
    if ($muggle) {
        # first see if the original uncensored images are around
        if (check_for_backups($params, \$fileArgs)) {
            # We're good to go. fileArgs has been edited to contain the paths for the original uncensored images
            print "Making stamps from backup images\n";
        } elsif (($options & $PSTAMP_REQUIRE_UNCENSORED) and ($stage ne 'chip')) {
            # user required uncensored but since stage isn't chip we can't rebuild them
            my_die("uncensored inputs not available for job $job_id", $job_id, $PSTAMP_NOT_AVAILABLE, 'stop');
        } elsif (($options & $PSTAMP_REQUEST_UNCENSORED) and ($params->{state} eq 'update') and ($stage ne 'chip')) {
            # we can only restore pixels for chip stage images if the data has been updated.
            # XXX: this test is not quite good enough. If all components have been updated then the
            # state will be 'full' But this will get us going.
            print "Run state is update: will make stamps from destreaked $stage images.\n";
            # make stamps from uncensored images
            $muggle = 0;
        } else {
            # Try and replace the streaks from the recovery images

            my $temp_dir = metadataLookupStr($ipprc->{_siteConfig}, "TEMP.DIR") or &my_die("Unable to find temporary directory in site configuration", $job_id, $PS_EXIT_CONFIG_ERROR);
            $tmproot = tempdir("$temp_dir/psjob.$job_id.XXXX", CLEANUP => !$save_temps);
            my $muggle_command = "$streaksreplace -stage $stage -tmproot $tmproot";
            # find the "directory" of the input path_base
            my $inputdir = dirname($image);
            my $base = basename($image);
            $tmpImage = "$tmproot/$base";

            @file_list = ();

            # XXX: We should get the recovery_path_base from the magicDSFile but that requires a bunch of file rule and
            # stage work. Import the rule here.
            my $recImage = "$inputdir/REC_$base";
            my $newFileArgs = " -file $tmpImage";
            $muggle_command .= " -image $image -recimage $recImage";
            push @file_list, $recImage;

            if ($mask) {
                $base = basename($mask);
                $tmpMask = "$tmproot/$base";
                $newFileArgs .= " -mask $tmpMask";
                my $recMask = "$inputdir/REC_$base";
                $muggle_command .= " -mask $mask -recmask $recMask";
                push @file_list, $recMask;
            }

            if ($variance) {
                $base = basename($variance);
                $tmpVariance = "$tmproot/$base";
                $newFileArgs .= " -variance $tmpVariance";
                my $recVariance = "$inputdir/REC_$base";
                $muggle_command .= " -weight $variance -recweight $recVariance";
                push @file_list, $recVariance;
            }
            if (check_files(0, @file_list)) {
                # recovery files exist and are accessible restore the excised pixels

                my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
                    run(command => $muggle_command, verbose => $verbose);
                unless ($success) {
                    my $exitStatus = WEXITSTATUS($error_code);
                    my_die( "streaksreplace failed with error code: $exitStatus", $job_id, $exitStatus);
                }

                # set the ppstamp file arguments
                $fileArgs = $newFileArgs;
            } else {
                if ($options & $PSTAMP_REQUIRE_UNCENSORED) {
                    my_die( "unable to restore uncensored images", $job_id, $PSTAMP_NOT_AVAILABLE);
                }
                # just make stamps from the censored images
                print "Unable to restore uncensored images, will extract stamps from censored images\n";
                # these files won't be used so zap them
                ($tmpImage, $tmpMask, $tmpVariance, $tmproot) = (undef, undef, undef, undef);
            }
        }
    }

    my $command = "$ppstamp $outputBase $argString $fileArgs";
    $command .= " -dbname $dbname" if $dbname;
    $command .= " -dbserver $dbserver" if $dbserver;
    $command .= " -stage $stage";
    $command .= " -no_censor_masked" unless $nan_masked;
    my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
        run(command => $command, verbose => $verbose);

    my $exitStatus;
    if (WIFEXITED($error_code)) {
        $exitStatus = WEXITSTATUS($error_code);
    } else {
        print STDERR "ppstamp failed error_code: $error_code\n";
        $exitStatus = $PS_EXIT_SYS_ERROR;
    }

    if ($exitStatus == 0) {
        my $reglist = "$outdir/reglist$job_id";

        my $F;
        open $F, ">$reglist" or my_die( "can't open $reglist for output", $job_id, $PS_EXIT_UNKNOWN_ERROR);

        # Figure out what output images were produced

        # Note: we are assuming the contents of the PSTAMP filerules here.
        my %extensions = ( $PSTAMP_SELECT_IMAGE  => "fits", 
                           $PSTAMP_SELECT_MASK   => "mk.fits",
                           $PSTAMP_SELECT_VARIANCE => "wt.fits");

        my $output_mask = $options & ($PSTAMP_SELECT_IMAGE | $PSTAMP_SELECT_MASK | $PSTAMP_SELECT_VARIANCE);

        foreach my $key (keys (%extensions)) {
            my $do_this_one = $key & $output_mask;

            next if (! $do_this_one);

            my $extension = $extensions{$key};

            my $basename = basename($outputBase);

            my $filename = "${basename}.${extension}";
            my $path  = "${outputBase}.${extension}";

            # XXX is pstamp always the right file type, if not where do we get the right one?
            print $F file_registration_line($filename, $path, "pstamp") . "\n";
        }

        get_other_outputs($F, $outputBase, $options, $params);

        close $F;
        $jobStatus = $PS_EXIT_SUCCESS;
    } elsif ($exitStatus == $PSTAMP_NO_OVERLAP) {
        $jobStatus = $PSTAMP_NO_OVERLAP;
    } else {
        my_die( "ppstamp failed with error code: $exitStatus", $job_id, $exitStatus);
    }
} elsif ($jobType eq "get_image") {

    my $uri = "";
    my $command = "$pstamp_get_image_job --job_id $job_id --output_base $outputBase --rownum $rownum";
    $command .= " --dbname $dbname" if $dbname;
    $command .= " --dbserver $dbserver" if $dbserver;
    $command .= " --verbose" if $verbose;
    my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
        run(command => $command, verbose => $verbose);

    if ($success) {
        $jobStatus = $PS_EXIT_SUCCESS;
    } else {
        $jobStatus = $error_code >> 8;
        my_die( "pstamp_get_image_job failed with error code: $jobStatus", $job_id, $jobStatus);
    }
} elsif ($jobType eq "detect_query") {
    # detect_query jobs are basically holders to note that we need a file updated before we can continue.
    # Load the argument list that dqueryparse should have created the first time it ran, so we know how to
    # run it again the same way.
    my $argslist = "$outputBase/parse.args";
    open ARGSLIST, "<$argslist" or my_die("failed to open argslist file $argslist", $job_id, $PS_EXIT_UNKNOWN_ERROR);
    my $argString = <ARGSLIST>;
    close ARGSLIST;
    chomp $argString;

    # XXX: should we do any other sanity checking?
    my_die("arglist file $argslist is empty", $job_id, $PS_EXIT_DATA_ERROR) if !$argString;

    my $command = "$dqueryparse --job_id $job_id $argString";
    $command .= " --dbname $dbname" if $dbname;
    $command .= " --dbserver $dbserver" if $dbserver;
    $command .= " --verbose" if $verbose;
    my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
	run(command => $command, verbose => $verbose);

    # Although dqueryparse can potentially force more updates (and more jobs), it should still return success.
    if ($success) {
	$jobStatus = $PS_EXIT_SUCCESS;
    } else {
	$jobStatus = $error_code >> 8;
	my_die("dqueryparse.pl failed with error code: $jobStatus", $job_id, $jobStatus);
    }
} else {
    my_die("unknown jobType $jobType found", $job_id, $PS_EXIT_PROG_ERROR);
}

# mark the job stopped in the database
{
    my $command = "$pstamptool -updatejob -job_id $job_id -set_state stop"; 
    $command .= " -set_fault $jobStatus" if $jobStatus;
    $command .= " -dbname $dbname" if $dbname;
    $command .= " -dbserver $dbserver" if $dbserver;
    if (!$no_update) {
        my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
            run(command => $command, verbose => $verbose);
        unless ($success) {
            die("Unable to perform $command: $error_code");
        }
    } else {
        print STDERR "skipping command: $command\n"
    }
}

exit 0;

# create a string to be passed as input to dsreg when registering this file in a fileset
# XXX move this to a module so it can be shared

sub file_registration_line {
    my $filename = shift;
    my $path     = shift;
    my $filetype = shift;
    if (-e $path) {
        my @finfo = stat($path);
        my_die("failed to stat $path", $job_id, $PS_EXIT_UNKNOWN_ERROR) unless (@finfo); 
        my $bytes = $finfo[7];
        my $md5sum = file_md5_hex($path);

        return "$filename|$bytes|$md5sum|$filetype|";
    } else {
        my_die("$filename not found at $path", $job_id, $PS_EXIT_UNKNOWN_ERROR);
    }
}
        
sub get_other_outputs {
    my $f = shift;
    my $output_base = shift;
    my $options = shift;
    my $params = shift;

    if ($options & ( $PSTAMP_SELECT_CMF | $PSTAMP_SELECT_PSF | $PSTAMP_SELECT_BACKMDL)) {
        if (!$params) {
            $params = read_params_file($output_base);
        }

        my $stage = $params->{stage};

        # raw files don't have any other data products
        return 1 if $stage eq "raw";

        # add the other data products if they are selected and exist.
        # silently skip them if they don't exist. Perhaps this should be
        # detected in pstampparse so that the user can be notified with 
        # a message in parse_error.txt ("warp does not have a background model")
        my $psf_file = $params->{psf} if ($options & $PSTAMP_SELECT_PSF);
        my $backmdl_file = $params->{backmdl} if ($options & $PSTAMP_SELECT_BACKMDL);
        my $pattern_file = $params->{pattern} if ($options & $PSTAMP_SELECT_BACKMDL);
        my $cmf_file;
        if ($stage ne 'chip') {
            # we don't ship chip stage cmf files because they may not be censored
            $cmf_file = $params->{cmf} if ($options & $PSTAMP_SELECT_CMF);
        }

        my $outdir = dirname($output_base);
        my $basename = basename($output_base);
        my ($rownum, $jobnum, $therest) = split /_/, $basename;
        &my_die("failed to split basename: $basename", $job_id, $PS_EXIT_CONFIG_ERROR) 
            if (!$therest or !$rownum or !$jobnum);

        my $prefix = "${rownum}_${jobnum}_";

        if ($cmf_file) {
            print "cmf file is $cmf_file\n";
            copy_and_register_file($f, $cmf_file, $outdir, $prefix);
        }
        if ($psf_file) {
            print "psf_file is $psf_file\n";
            copy_and_register_file($f, $psf_file, $outdir, $prefix);
        }
        if ($backmdl_file) {
            print "backmdl_file is $backmdl_file\n";;
            copy_and_register_file($f, $backmdl_file, $outdir, $prefix);
        }
       if (0) {
        # don't enable this yet
        if ($pattern_file) {
            print "pattern_file is $pattern_file\n";;
            copy_and_register_file($f, $pattern_file, $outdir, $prefix);
        }
       }
    }
}

sub read_params_file {
    my $output_base = shift;

    my $params_file = $output_base . ".mdc";
    open (IN, "<$params_file") 
        or my_die("failed to open params file: $params_file", $job_id, $PS_EXIT_UNKNOWN_ERROR);

    my $data = $mdcParser->parse(join "", (<IN>))
        or my_die("failed to parse params file: $params_file", $job_id, $PS_EXIT_UNKNOWN_ERROR);

    my $components = parse_md_list($data);

    my $n = scalar @$components;
    if ($n != 1) {
        my_die("params file $params_file contains unexpected number of components: $n",
                $job_id, $PS_EXIT_PROG_ERROR);
    }
    return $components->[0];
}

# copy_and_register_file ($f, $src, $destdir, $prefix);
sub copy_and_register_file {
    my $F = shift;
    my $src = shift;
    my $destdir = shift;
    my $prefix = shift;

    my $fn = $prefix . basename($src);
    my $dst = "$destdir/$fn";

    my $resolved = $ipprc->file_resolve($src);

    my_die("failed to resolve $src", $job_id, $PS_EXIT_UNKNOWN_ERROR) if !$resolved;

    if (!-e $resolved) {
        print STDERR "$src does not exist, skipping\n";
        return;
    }


    copy($resolved, $dst) or my_die("failed to copy $resolved to $dst", $job_id, $PS_EXIT_UNKNOWN_ERROR);

    print $F file_registration_line($fn, $dst, "fits") . "\n";
}

sub check_files {
    my $error_code = shift;
    my $return_code = 1;
    foreach my $f (@_) {
        if (!$ipprc->file_exists($f)) {
            if ($error_code) {
                my_die( "file $f does not exist:", $job_id, $error_code, 'stop');
            } else {
                print STDERR "file $f does not exist\n";
                $return_code = 0;
            }
        }
    }
    return $return_code;
}

sub check_for_backups {
    my $params = shift;
    my $r_fileArgs = shift;

    my $command = "$magicdstool -destreakedfile -stage $params->{stage} -stage_id $params->{stage_id} -component $params->{component} -dbname $params->{imagedb}";
    my $results = runToolAndParse($command, $verbose);
    my $dsComponent  = $results->[0];

    if ($dsComponent and ($dsComponent->{data_state} eq 'full')) {

        print "magicDSFile state is full. Backup images should exist\n";

        # replace the file names with the backup paths
        # fileArgs has the form:  -file imagename [-mask maskname] [-weight weightname]
        my @args = split " ", $$r_fileArgs;

        my $newFileArgs;
        while (@args) {
            my $a = shift @args;
            my $f = shift @args;
            my_die( "unexpected fileArg list $$r_fileArgs", $job_id, $PS_EXIT_PROG_ERROR, 'run') if !$a or !$f;

            my $backup_image_name = $ipprc->destreaked_filename($f);
            my_die( "failed to extract backup image name from $f", $job_id, $PS_EXIT_PROG_ERROR, 'run') if !$backup_image_name;
            $newFileArgs .= " $a $backup_image_name";
        }
        $$r_fileArgs = $newFileArgs;

        return 1;
    }

    return 0;
}

sub my_die
{
    my $msg = shift;            # Warning message on die
    my $job_id = shift;         # job identifier
    my $exit_code = shift;      # Exit code to add
    my $job_state = shift;      # new pstampJob.state

    $exit_code = $PS_EXIT_PROG_ERROR unless $exit_code;
    $job_state = 'run' unless $job_state;

    carp($msg);
    if (defined $job_id and not $no_update) {
        my $command = "$pstamptool -updatejob";
        $command .= " -job_id $job_id";
        $command .= " -set_fault $exit_code";
        # XXX: fix pstamptool to not require -state when -fault with nonzero value is provided
        $command .= " -set_state $job_state";
        $command .= " -dbname $dbname" if defined $dbname;
        $command .= " -dbserver $dbserver" if defined $dbserver;
        system($command);
    }
    exit $exit_code;
}

