#!/usr/bin/perl -w
use strict;
use Getopt::Std qw( getopts );
use POSIX qw( );

( my $Script = $0 ) =~ s{^.*/}{};

my $time = time;
my $date = POSIX::strftime( "%F %T", localtime $time );

$ENV{SGE_ROOT} or die "Error $Script: \$SGE_ROOT is not set\n";
$ENV{SGE_CELL} ||= "default";    # sge_cell

my ($siteLocation) = join "/" => ( $ENV{SGE_ROOT}, $ENV{SGE_CELL}, "site" );

my $config = {
    -site   => ( $siteLocation || '' ),
    -cache  => ( $siteLocation || '' ) . "/cache/" . "qlicserver.xml",
    -limits => ( $siteLocation || '' ) . "/" . "qlicserver.limits",
};

# ------------------------------------------------------------------------
# utils

#
# calculate age from an epoch value
#
sub age {
    my ( $a, $b ) = @_;
    my $diff = ( $a - $b );

    my $sign = '';
    if ( $diff < 0 ) {    # handle negatives
        $sign = '-';
        $diff = abs($diff);
    }

    sprintf "$sign%d:%02d:%02d",    # format into hh:mm:ss
      ( int $diff / 3_600 ),        # hh
      ( ( int $diff / 60 ) % 60 ),  # mm
      ( $diff % 60 );               # ss
}

#
# change hash references to a comma-delimited string of key=value entries
#
sub hashRefToString {
    map {
        my $ref = $_;
        ref $ref ? map { "$_=$ref->{$_}" } sort keys %$ref : ()
    } @_;
}

#
# extract attrib="value" ... attrib="value"
#
sub parseXMLattrib {
    my $str = shift || '';
    my %attr;

    while ($str =~ s{^\s*(\w+)=\"([^\"]*)\"}{}s
        or $str =~ s{^\s*(\w+)=\'([^\']*)\'}{}s )
    {
        $attr{$1} = $2;
    }

    %attr;
}

# ------------------------------------------------------------------------
# processing of qlicserver.xml output
#
package qlicCache;

sub new {
    my $caller = shift;
    my $file   = shift;

    my $self = bless {
        ( map { $_ => {} } qw( -rc ) ),
        ( map { $_ => '' } qw( -date -host -age ) )
    }, $caller;

    -f $file or return $self;

    local $_ = do { local ( @ARGV, $/ ) = $file; <>; };

    # strip out all xml comments
    s{<!--.*?-->\s*}{}sg;

    # get the header section
    s{^(.+?)</query>}{}s or return;
    my $header = $1;

    for ($header) {
        if (m{<host (?:\s+([^<>]*))? > \s*(.+?)\s* </host>}sx) {
            my ( $attr, $content ) = ( $1, $2 );
            $self->{-host} = $content;
        }
        if (m{<time (?:\s+([^<>]*))? > \s*(.+?)\s* </time>}sx) {
            my ( $attr, $content ) = ( $1, $2 );
            my %attr = main::parseXMLattrib($attr);

            $self->{-age} = main::age( $time, $attr{epoch} ) if $attr{epoch};
            $self->{-date} = $content;
        }
    }

    # only retain the <resources> contents
    s{^.*?<resources>}{}s and s{</resources>.*$}{}s or return;

    ## process <resource .../>
    while (s{<resource (?:\s+([^/<>]*))? />}{}sx) {
        my ($attr) = ($1);
        my %attr   = main::parseXMLattrib($attr);
        my $name   = delete $attr{name};

        if ( defined $name ) {
            for ( keys %attr ) {
                $self->{-rc}{$name}{$_} = $attr{$_} || 0;
            }
        }
    }

    ## process <resource ...> CONTENT </resource>
    while (s{<resource (?:\s+([^<>]*))? > \s* (.*?) \s* </resource>}{}sx) {
        my ( $attr, $content ) = ( $1, $2 );
        my %attr = main::parseXMLattrib($attr);
        my $name = delete $attr{name};

        if ( defined $name ) {
            for ( keys %attr ) {
                $self->{-rc}{$name}{$_} = $attr{$_} || 0;
            }
        }

        ## process <user ...> COUNT </user>
        while ( $content =~ s{<user (?:\s+([^<>]*))? >\s*(\d+)\s*</user>}{}sx )
        {
            my ( $attr, $count ) = ( $1, $2 );
            my %attr = main::parseXMLattrib($attr);
            my $user = delete $attr{name};
            my $host = delete $attr{host};

            if ( defined $user ) {
                if ( defined $host ) {
                    $self->{-rc}{$name}{-where}{"$user\@$host"} = $count;
                }
                else {
                    ## tag waiting with '[]'
                    $self->{-rc}{$name}{-where}{$user} = "[$count]";
                }
            }
        }
    }

    return $self;
}

sub query {
    my $self = shift;
    %{ $self->{-rc} };
}

sub dump {
    use Data::Dumper;
    print Dumper( $_[0] );
}

sub available {
    my $self = shift;
    sort keys %{ $self->{-rc} };
}

# ------------------------------------------------------------------------
# processing of qlicserver.limits FILE or DIRECTORY
#
package qlicLimits;

sub new {
    my $caller     = shift;
    my $diskValues = shift;

    my $self = bless {}, $caller;

    my $fileString;
    if ( defined $diskValues and -f $diskValues ) {
        $fileString = do {
            local ( *FILE, $/ );

            if ( open FILE, $diskValues ) {
                <FILE>;
            }
            else {
                undef;
            }
        };
    }
    elsif ( defined $diskValues and -d $diskValues ) {
        local *DIR;
        my $dir = $diskValues;
        if ( opendir DIR, $dir ) {
            my @files = grep { -f "$dir/$_" and -r _ } readdir DIR;
            for my $name (@files) {
                my $limit;

                # use the last value
                if ( open FILE, "$dir/$name" ) {
                    $limit = ( map { /^\s*(-?\d+)\s*$/ } <FILE> )[-1];
                }
                defined $limit or $limit = "NONE";
                $self->{$name} = $limit;
            }
        }
    }

    for ($fileString) {
        defined or next;

        # strip out all comments
        s{<!--.*?-->\s*}{}sg;

        ## process <resource .../> and <resource ...></resource>
        while (s{<resource (?:\s+([^/<>]*))? />}{}sx
            or s{<resource (?:\s+([^/<>]*))? >\s*</resource>}{}sx )
        {
            my %attr  = main::parseXMLattrib($1);
            my $name  = delete $attr{name};
            my $limit = delete $attr{limit};
            if ( defined $name and defined $limit ) {
                $self->{$name} = $limit;
            }
        }
    }

    return $self;
}

sub change {
    my $caller     = shift;
    my $diskValues = shift;
    my @args       = @_;

    @args or return;

    my ( %pending, %adjusted );

    for (@args) {
        s{,}{ }g;    # comma -> space-delimited

        my %h = map { /^(.+?)=(.*)$/ } split;
        for ( keys %h ) {
            defined $h{$_} and length $h{$_} or $h{$_} = "NONE";
            $pending{$_} = $h{$_};
        }
    }

    keys %pending or return;

    my $user = getpwuid $<;
    if ( defined $diskValues and -d $diskValues ) {
        local *DIR;
        my $dir = $diskValues;
        if ( opendir DIR, $dir ) {
            my @files = grep { -f "$dir/$_" and -w _ } readdir DIR;

            for my $name (@files) {
                if ( exists $pending{$name} ) {
                    local *FILE;
                    if ( open FILE, ">", "$dir/$name" ) {
                        print FILE "# adjusted by $user $date\n",
                          "$pending{$name}\n";
                        $adjusted{$name} = delete $pending{$name};
                    }
                }
            }
        }
    }

    for ( [ "adjusted" => \%adjusted ], [ "not adjusted" => \%pending ], ) {
        my ( $label, $href ) = @$_;
        keys %$href or next;
        print "$label: ",
          join( "," => map { qq{$_=$href->{$_}} } sort keys %$href ), "\n";
    }
}

sub output {
    my $self = shift;

    my @list = map { qq{$_=$self->{$_}} } sort keys %$self;

    print "limits: ";
    if (@list) {
        print join( "," => @list ), "\n";
    }
    else {
        print "NONE\n";
    }
}

# ------------------------------------------------------------------------
# main
#

package main;
import qlicCache;

# ------------------------------------------------------------------------
sub usage {
    $! = 0;    # clean exit
    warn "@_\n" if @_;
    die <<"USAGE";
usage:
    $Script [OPTION]
    $Script [OPTION] resource=limit .. resource=limit

with options:
  -c FILE  alternative location for the license cache
  -C FILE  alternative location for the license limit
  -d       dump cache file as raw xml
  -D       dump license cache in perl format
  -f       display free licenses only
  -l       list license limit
  -q       display free licenses via qhost query
  -u       license usage via 'lacct'
  -U       license usage per user via 'lacct -u'
  -w       show who/where ('[]' indicates waiting jobs)
  -h       this help

* extract / display information for the GridEngine license cache
  $config->{-cache}

* adjust / display information for the license limits
  $config->{-limits}

copyright (c) 2003-10 <Mark.Olesen\@faurecia.com>

Licensed and distributed under the Creative Commons
Attribution-NonCommercial-ShareAlike 3.0 License.
http://creativecommons.org/licenses/by-nc-sa/3.0
USAGE
}

# -------------------------------------------------------------------------
my %opt;
getopts( 'hc:C:DdflqUuw', \%opt ) or usage();
$opt{h} and usage();

if ( $opt{U} ) {
    ## user-based usage from accounting
    my @query = ( "lacct", "-u" );
    exec @query;
    exit 2;
}
elsif ( $opt{u} ) {
    ## usage from accounting
    my @query = ( "lacct", $opt{w} ? "-w" : () );
    exec @query;
    exit 2;
}

if ( $opt{d} ) {
    my $file = $opt{c} || $config->{-cache};

    if ( not -f $file and $file !~ m{/} ) {
        $file = "$config->{-site}/$file";
    }

    -f $file or exit 1;

    local $_ = do { local ( @ARGV, $/ ) = $file; <>; };

    # do a basic check for well-formed xml
    # this might not be the case if there is a race condition
    # and the file has not been fully written

    # check for '<?xml version="1.0"?><someTag ...' at the start
    # and '</someTag>' at the end
    my ($tag) = m{^<\?xml (?:\s+[^<>]*?)? \?>\s*\<(\w+)}sx;
    unless ( $tag and m{</\Q$tag\E>\s*$} ) {
        sleep 2;    ## wait and try again
        $_ = do { local ( @ARGV, $/ ) = $file; <>; };
    }

    $_ ||= '';      ## avoid uninitialized

    print;
    exit;
}

if ( $opt{q} ) {
    my %free =
      map  { /^\s+gc:(\S+?)=(\d\S*)\s*$/ }
      grep { /^global/ ... /^\S/ } qx{qhost -F};

    $_ += 0 for values %free;

    for ( sort keys %free ) {
        my $intval = $free{$_} + 0;    # normalize as integers
        print "$_=$intval\n";
    }
    exit;
}

if ( $opt{l} ) {
    qlicLimits->new( $opt{C} || $config->{-limits} )->output();
    exit;
}

if (@ARGV) {
    qlicLimits->change( $opt{C} || $config->{-limits}, @ARGV );
    exit;
}

my $info = qlicCache->new( $opt{c} || $config->{-cache} );
$info->dump() if $opt{D};

my %resources = $info->query();

#
# display free licenses
#
if ( $opt{f} ) {
    for my $resource ( sort keys %resources ) {
        my $count = $resources{$resource} || {};
        my $free  = $count->{free}        || 0;
        print "$resource=$free\n" if $free;
    }

    exit;
}

#
# display who/where
#
if ( $opt{w} ) {
    my ($len) = sort { $b <=> $a } map { length } keys %resources;
    $len += 2;

    my $fmt    = "%-${len}s";
    my $indent = sprintf "\n$fmt", '';

    for my $resource ( sort keys %resources ) {
        my $count = $resources{$resource} || {};
        my @list = hashRefToString( $count->{-where} );
        if (@list) {
            printf $fmt, $resource;
            print join( $indent => @list ), "\n";
        }
    }

    exit;
}

#
# define table output format
#
my @outputList;

format =
# name      total limit extern intern wait free
@<<<<<<<<<<< @>>>>> @>>>>> @>>>>> @>>>>> @>>>> @>>>>
@outputList
.

#
# display table header
# --------------------
{
    my @info = qw( host age );
    my ($len) = sort { $b <=> $a } map { length } @info;

    print map {
        my $k = sprintf "%-${len}s", $_;
        my $v = $info->{"-$_"};
        $v ? "$k = $v\n" : ();
    } @info;

    print "\n";

    @outputList = qw( resource total limit extern intern wait free );

    write;
    s/./-/g for @outputList;
    write;
}

#
# display table body
# ------------------
for my $resource ( sort keys %resources ) {
    my $count = $resources{$resource} || {};
    @outputList =
      ( $resource, @{$count}{qw( total limit extern intern waiting free )} );

    my $type = $resources{$resource}{type} || 0;

    # no total?
    $_ ||= "?" for $outputList[1];

    if ( $type eq "intern" ) {
        $_ = "*" for $outputList[3];
    }
    elsif ( $type eq "track" ) {
        $_ = "*" for $outputList[4];
    }

    # easy to read representation for zero
    for (@outputList) {
        defined $_ or $_ = '.';
    }

    write;
}

# ------------------------------------------------------------------------