#!/usr/bin/env perl

use Carp;
use warnings;
use strict;

## report the program and machine
use Sys::Hostname;
my $host = hostname();
print "\n\n";
print "Starting script $0 on $host\n\n";

use vars qw( $VERSION );
$VERSION = '0.01';

use IPC::Cmd 0.36 qw( can_run run );
use File::Temp qw( tempfile );
use File::Basename qw( basename );
use Digest::MD5::File qw( file_md5_hex );
use PS::IPP::Metadata::Config;
use PS::IPP::Metadata::List qw( parse_md_list );

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

my $ipprc = PS::IPP::Config->new(); # IPP configuration

use Getopt::Long qw( GetOptions :config auto_help auto_version gnu_getopt );
use Pod::Usage qw( pod2usage );

# filerules that are not included in 'cleaned' distribution bundles
# I could simplify the function below by using a two level hash but
# that would be less clear I think
my %chip_cleaned = ( 'PPIMAGE.CHIP' => 'image',
                     'PPIMAGE.CHIP.MASK' => 'mask',
                     'PPIMAGE.CHIP.VARIANCE' => 'variance' );
my %camera_cleaned = ( 'PSASTRO.OUTPUT.MASK' => 'mask' );
my %fake_cleaned;
my %warp_cleaned = ( 'PSWARP.OUTPUT' => 'image',
                     'PSWARP.OUTPUT.MASK' => 'mask',
                     'PSWARP.OUTPUT.VARIANCE' => 'variance' );
my %diff_cleaned = ( 'PPSUB.OUTPUT' => 'image',
                     'PPSUB.OUTPUT.MASK' => 'mask',
                     'PPSUB.OUTPUT.VARIANCE' => 'variance' );
my %stack_cleaned = ( 'PPSTACK.OUTPUT' => 'image',
                      'PPSTACK.OUTPUT.MASK' => 'mask',
                      'PPSTACK.OUTPUT.VARIANCE' => 'variance' );


# Look for programs we need
my $missing_tools;
my $disttool   = can_run('disttool') or (warn "Can't find disttool" and $missing_tools = 1);
my $streaksrelease   = can_run('streaksrelease') or (warn "Can't find streaksrelease" and $missing_tools = 1);
if ($missing_tools) {
    warn("Can't find required tools.");
    exit($PS_EXIT_CONFIG_ERROR);
}

# Parse the command-line arguments
my ($dist_id, $camera, $stage, $stage_id, $component, $path_base, $chip_path_base, $clean);
my ($outdir, $run_state, $data_state, $magicked, $no_magic, $poor_quality);
my ($dbname, $save_temps, $verbose, $no_update, $logfile);

GetOptions(
           'dist_id=s'      => \$dist_id,    # distribution run identifier
           'camera=s'       => \$camera,     # camera for evaluating file rules
           'stage=s'        => \$stage,      # raw, chip, warp, or diff
           'stage_id=s'     => \$stage_id,   # exp_id, chip_id, warp_id, or diff_id
           'component=s'    => \$component,  # the class_id or skycell_id
           'path_base=s'    => \$path_base,  # path_base of the input
           'chip_path_base=s'=> \$chip_path_base,  # path base for camera stage (to enable us to find the mask filefor raw images)
           'state=s'        => \$run_state,  # state of the run
           'data_state=s'   => \$data_state, # data_state for this component
           'poor_quality'   => \$poor_quality,  # the processing for this component did not produced images
           'no_magic'       => \$no_magic,   # magic is not required for this distribution run
           'magicked'       => \$magicked,   # magicked state for this component
           'outdir=s'       => \$outdir,     # "directory" for outputs
           'clean'          => \$clean,      # create clean distribution
           'save-temps'     => \$save_temps, # Save temporary files?
           'dbname=s'       => \$dbname,     # Database name
           'verbose'        => \$verbose,    # Print stuff?
           'no-update'      => \$no_update,  # Don't update the database?
           'logfile=s'      => \$logfile,
           ) or pod2usage( 2 );

pod2usage( -msg => "Unknown option: @ARGV", -exitval => 2 ) if @ARGV;
pod2usage( -msg => "Required options: --dist_id --camera --stage --stage_id --component --path_base --outdir",
           -exitval => 3) unless
    defined $dist_id and
    defined $camera and
    defined $stage and
    defined $stage_id and
    defined $component and
    defined $path_base and
    defined $outdir;

$ipprc->redirect_output($logfile) if $logfile;

if (($stage eq 'raw') and !$clean and !defined $chip_path_base) {
    pod2usage( -msg => "Required options: --chip_path_base for raw stage", -exitval => 3);
}

$ipprc->define_camera($camera);

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

# making a clean bundle of raw images doesn't make sense

# create the output directories if it is not a nebulous path and it doesn't exist
if (index($outdir, "neb://") != 0) {
    if (! -e $outdir ) {
        my $code = system "mkdir -p $outdir";
        &my_die("cannot create output directory $outdir", $dist_id, $component,
                $code >> 8) if $code;
    }
}

# Get the list of data products for this component
# note: We my_die in get_file_list if something goes wrong. 

my $file_list = get_file_list($stage, $component, $path_base, $clean);

if (($stage ne 'raw') and ($stage ne 'fake')) {
    # If the file list is empty it is an error because we should at least get a config dump file
    # except for fake stage which doesn't do anything yet
    &my_die("empty file list", $dist_id, $component, $PS_EXIT_CONFIG_ERROR) if (!scalar @$file_list);
}

# set up directory for temporary files

my $tmpdir  = "$outdir/tmpdir.$dist_id.$component";
if (-e $tmpdir) {
    if (-d $tmpdir) {
        my $rc = system "rm -r $tmpdir";
        &my_die("cannot rm $tmpdir return code: $rc", $dist_id, $component, $PS_EXIT_UNKNOWN_ERROR) if $rc;
    } else {
        unlink $tmpdir or &my_die("cannot delete $tmpdir", $dist_id, $component, $PS_EXIT_UNKNOWN_ERROR);
    }
}
mkdir $tmpdir or &my_die("cannot create temporary directory $tmpdir", $dist_id, $component,
                       $PS_EXIT_UNKNOWN_ERROR);

#
# we need to run set masked pixels to NAN in the image and variance images
# unless
#   1. we are building a clean distribution bundle
#   2. magic is not required for this distRun
#   3. the processing for the component produced no images (warp or diff with bad quality for example)
my $nan_masked_pixels = ! ($clean || (($stage eq "camera") || ($stage eq 'fake') || ($stage eq 'stack')) || $no_magic || $poor_quality);

my ($image, $mask, $variance);

# foreach my $file_rule (keys %$file_list) {
my $num_files = 0;
foreach my $file (@$file_list) {
    # check whether this file rule refers to an image, mask, or variance fits image
    my $file_rule = $file->{file_rule};
    my $image_type = get_image_type($stage, $file_rule);

    # if we are building a clean bundle skip this rule
    next if $clean && $image_type;

    # if magic is required, don't ship jpegs or binned fits images
    next if !$no_magic && (($file_rule =~ /.BIN1/) or ($file_rule =~ /.BIN2/) or ($file_rule =~ /.JPEG1/)
            or ($file_rule =~ /.JPEG2/));

    my $file_name = $file->{name};
    my $base = basename($file_name);
    my $path = $ipprc->file_resolve($file_name);

    if (!$path) {
        # skip this file if $poor_quality
        # this is for compatability with older runs which don't have the files list in the
        # config dump.
        # Once we give up on supporting that we can remove the next line. (If the file is in the list
        # it must exist)
        next if $poor_quality;

        &my_die("failed to resolve  $file_name", $dist_id, $component, $PS_EXIT_DATA_ERROR);
    }

    # open the file to make sure it exists (and to work around the failed mount phenomena)
    my $fh = open_with_retries($path);
    close $fh;

    # we need to pre-process the image before adding to the bundle. Save the path names.
    # the images will be created below
    $num_files++;
    if ($image_type && $nan_masked_pixels) {
        # save the 
        if ($image_type eq 'image') {
            $image = $path
        } elsif ($image_type eq 'mask') {
            $mask = $path;
        } elsif ($image_type eq 'variance') {
            $variance = $path;
        } else {
            &my_die("invalid image type found: $image_type", $dist_id, $component,
                       $PS_EXIT_PROG_ERROR);
        }
    } else {
        # create a symbolic link from the file in the nebulous repository
        # in the temporary directory
        symlink $path, "$tmpdir/$base";
    }
}

if ($nan_masked_pixels) {
    # One last check as to whether magic has been applied to the inputs

    # Note: the sql for disttool -pendingcomponent won't select a component that 
    # requires magic and hasn't been magicked, but we check again here

    if (!($magicked or $no_magic)) {
        &my_die("cannot create distribution bundle ${stage}_id $stage_id because the data has not been magic desreaked", $dist_id, $component, $PS_EXIT_DATA_ERROR);
    }

    &my_die("no image found in file list", $dist_id, $component, $PS_EXIT_CONFIG_ERROR) if !$image;
    &my_die("no mask image found in file list", $dist_id, $component, $PS_EXIT_CONFIG_ERROR) if !$mask;
    &my_die("no variance image found in file list", $dist_id, $component, $PS_EXIT_CONFIG_ERROR) if !$variance;

    my $class_id;
    # run streaksrelease to set masked pixels to NAN
    if ($stage eq "raw") {
        $class_id = $component;
        # we can use the chip mask because disttool demands that magic have been run
        # and so the camera mask and the chip mask are the same
        $mask = $ipprc->filename("PPIMAGE.CHIP.MASK", $chip_path_base, $component);
        my $fh = open_with_retries($mask);
        close $fh;
    } elsif ($stage eq "chip") {
        $class_id = $component;
    }

    my $command = "$streaksrelease -stage $stage -image $image -outroot $tmpdir";
    $command .= " -class_id $class_id" if $class_id;
    $command .= " -mask $mask" if $mask;
    $command .= " -chip_mask $mask" if ($stage eq 'chip' and $mask);
    $command .= " -weight $variance" if $variance;
    my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
        run(command => $command, verbose => $verbose);
    unless ($success) {
        $error_code = (($error_code >> 8) or $PS_EXIT_PROG_ERROR);
        &my_die("Unable to perform $command: $error_code", $dist_id, $component, $error_code);
    }
}


my $file_name;
my $bytes;
my $md5sum;

if ($num_files) {
    # create the tarfile
    # XXX TODO: create a file rule for the tar file name
    my $tbase = basename($path_base);
    $tbase .= ".$component" if $component;
    $file_name = "$tbase.tgz";
    my $tarfile = "$outdir/$file_name";

    my $command = "tar -C $tmpdir --owner=ipp --group=users -czhf $tarfile .";

    my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
        run(command => $command, verbose => $verbose);
    unless ($success) {
        $error_code = (($error_code >> 8) or $PS_EXIT_PROG_ERROR);
        &my_die("Unable to perform $command: $error_code", $dist_id, $component, $error_code);
    }

    # tell the module not to die on error
    $Digest::MD5::File::NOFATALS = 1;
    $md5sum = file_md5_hex($tarfile);
    &my_die("unable to compute md5sum for $tarfile", $dist_id, $component, $PS_EXIT_UNKNOWN_ERROR) if !$md5sum;

    $bytes = -s $tarfile;

    delete_tmpdir($tmpdir);
} else {
    # no files for this component
    $file_name = "none";
    $bytes = 0;
    $md5sum = "0";
}
{
    my $command = "$disttool -addprocessedcomponent -dist_id $dist_id -component $component";
    $command .= " -name $file_name -bytes $bytes -md5sum $md5sum";
    $command .= " -dbname $dbname" if defined $dbname;

    unless (defined $no_update) {
        my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
            run(command => $command, verbose => $verbose);
        unless ($success) {
            $error_code = (($error_code >> 8) or $PS_EXIT_PROG_ERROR);
            &my_die("Unable to perform $command: $error_code", $dist_id, $component, $error_code);
        }
    } else {
        print "skipping command $command\n";
    }
}

exit 0;

### Pau.

# return the image type (image, mask, or variance) if this file rule refers to 
# one of the big fits files that are not included in a clean distribution
sub get_image_type {
    my $stage = shift;
    my $rule = shift;
    my $type;

    if ($stage eq "raw") {
        if ($rule eq "RAW.IMAGE") {
            $type = 'image';
        }
    } elsif ($stage eq "chip") {
        $type = $chip_cleaned{$rule};
    } elsif ($stage eq "camera") {
        $type = $camera_cleaned{$rule};
    } elsif ($stage eq "fake") {
        $type = $fake_cleaned{$rule};
    } elsif ($stage eq "warp") {
        $type = $warp_cleaned{$rule};
    } elsif ($stage eq "diff") {
        $type = $diff_cleaned{$rule};
    } elsif ($stage eq "stack") {
        $type = $stack_cleaned{$rule};
    } else {
        &my_die("$stage is not a valid stage", $dist_id, $component, $PS_EXIT_CONFIG_ERROR);
    }
    return $type;
}

sub open_with_retries {
    my $name = shift;
    
    my $tries = 1;
    my $max_tries = 5;
    my $opened = 0;
    while (!$opened && ($tries <= $max_tries)) {
        $opened = open(IN, "<$name");
        if (!$opened) {
            print STDERR "WARNING failed to open $name re-try $tries\n";
            $tries++;
            sleep 5;
        }
    }
    
    &my_die("failed to open $name after $max_tries tries\n", $dist_id, $component,
                    $PS_EXIT_DATA_ERROR) if (!$opened);

    return *IN;
}

sub get_file_list {
    my $stage = shift;
    my $component = shift;
    my $path_base = shift;
    my $clean = shift;

    my @file_list;
    if ($stage eq "raw") {
        # XXX: TODO for now disttool sets path_base is set to the uri of the file
        # eventually rawImfile will have a path_base that we'll need to add '.fits'
        # XXX do we want to distribute the registration log files?
        if (!$clean) {
            my %file;
            $file{file_rule} = "RAW.IMAGE";
            $file{name} = $path_base;
            push @file_list, \%file;
        }
        return \@file_list;
    }

    # we get the list of output data products for this run from the config dump file that
    # is created when the run is done
    my $config_file_rule;
    if ($stage eq "chip") {
        $config_file_rule = "PPIMAGE.CONFIG";
    } elsif ($stage eq "camera") {
        $config_file_rule = "PSASTRO.CONFIG";
    } elsif ($stage eq 'fake') {
        # XXX: fake is a no op now return an emtpy list
        return \@file_list;
    } elsif ($stage eq "warp") {
        $config_file_rule = "PSWARP.CONFIG";
    } elsif ($stage eq "diff") {
        $config_file_rule = "PPSUB.CONFIG";
    } elsif ($stage eq "stack") {
        $config_file_rule = "PPSTACK.CONFIG";
    } else {
        &my_die("$stage is not a valid stage", $dist_id, $component, $PS_EXIT_CONFIG_ERROR);
    }
    my $config_file = $ipprc->filename($config_file_rule, $path_base, $component) or
                &my_die("can't get filename for config dump file: $config_file_rule", $dist_id, $component,
                    $PS_EXIT_CONFIG_ERROR);

    my %config_file_hash;

    # add the configuration file to the list
    $config_file_hash{file_rule} = $config_file_rule;
    $config_file_hash{name} = $config_file;
    push @file_list, \%config_file_hash;

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

    &my_die("failed to resolve name of config dump file: $config_file_rule", $dist_id, $component,
                    $PS_EXIT_CONFIG_ERROR) if (!$resolved);

    # we don't use the mdc parser because the perl parser is way is too slow for complicated config
    # files like this
    my $in = open_with_retries($resolved);

    my $line;
    while ($line = <$in>) {
        if ($line =~ /FILES.OUTPUT/) {
            # print "found FILES.OUTPUT\n";
            last;
        }
    }
    &my_die("config dump file does not contain FILES.OUTPUT: $config_file", $dist_id, $component,
                    $PS_EXIT_CONFIG_ERROR) if (!$line);

    while ($line = <$in>) {
        chomp $line;
        my ($key, $type, $val) = split " ", $line;
        # skip blank lines
        next if !$key;
        # we're done when we find END
        last if $key eq "END";
        # skip multi and other lines
        next if ($type ne "STR");

        # printf "%-32.32s   %s\n", $key, $val;

        &my_die("no value found for file rule $key in $resolved", $dist_id, $component,
                $PS_EXIT_CONFIG_ERROR) if (!$val);

        my %file;
        $file{file_rule} = $key;
        $file{name} = $val;
        push @file_list, \%file;
    }
    close $in;

    return \@file_list;
}




sub file_check
{
    my $file = shift;           # Name of file
    &my_die("Unable to find output file: $file", $dist_id, $component, $PS_EXIT_SYS_ERROR) unless
        $ipprc->file_exists($file);
}

sub my_die
{
    my $msg = shift;            # Warning message on die
    my $dist_id = shift;    # Magic DS identifier
    my $component = shift;      # class_id or skycell_id
    my $exit_code = shift;      # Exit code to add

    $exit_code = $PS_EXIT_PROG_ERROR unless defined $exit_code;

    my $command = "$disttool -addprocessedcomponent";
    $command   .= " -dist_id $dist_id";
    $command   .= " -component $component";
    $command   .= " -fault $exit_code";
    $command   .= " -dbname $dbname" if defined $dbname;

    # Add the processed file to the database
    unless ($no_update) {
        my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
            run(command => $command, verbose => $verbose);
        unless ($success) {
            carp("failed to update database for $dist_id $component");
        }
    } else {
        print "Skipping command: $command\n";
    }

    delete_tmpdir();

    carp($msg);
    exit $exit_code;
}

sub delete_tmpdir
{
    if (!$save_temps and $tmpdir and -e $tmpdir) {
        system "rm -r $tmpdir";
    }
}

__END__
