#! /usr/bin/env perl

use strict;
use warnings;
#use experimental "switch";

# requires File::Slurp. 
# In Debian derivatives
# apt install libfile-slurp-perl

# apt install libxml-libxml-perl libyaml-tiny-perl


BEGIN {
   use FindBin;
   use File::Spec;
   # use libraries from the directory this script is in
   use lib File::Spec->catdir($FindBin::Bin);
}

use Data::Dumper;
use YAML::Tiny;

# global variables
my $scriptDir = $FindBin::RealBin;
my $scriptName = $FindBin::Script;
my $confDir = "$scriptDir/conf";
my $dbDir = "$scriptDir/var";
my $nodeDBName = "$dbDir/node.yaml";
my $domainDBName = "$dbDir/domains.yaml";
my $nodePopulationDBName = "$dbDir/node_population.yaml";

# these contain the values from the databases
# loaded on demand
my $nodeDB;
my $virtDB;
my $nodePopulations;


my $dryRun = 1;
my $DEBUG = 0;

sub help {
   print "$0 command [argument]\n";
   print "where command is one of\n";
   print "\tnode update [node] [node]... # update a given node (or ALL)\n";
   print "\tnode list # display tab delimited list of node specs\n";
   print "\tnode scan # find domains on all nodes\n ";
   print "\tdomain update ALL|RUNNING|[domain] [domain]... # update domains\n";
   print "\tdomain list ALL|RUNNING|[domain] [domain]... # display tab delimited list of domain specs\n";
   print "\tcluster status # report of memory and vcpu status on all nodes\n";
}

sub readDB {
   my ($filename) = @_;
   my $yaml = YAML::Tiny->new( {} );
   if ( -f $filename ) {
      $yaml = YAML::Tiny->read( $filename );
   }
   return $yaml->[0];
}

sub writeDB {
   my ($filename,$data) = @_;
   my $yaml = YAML::Tiny->new( $data );
   $yaml->write( $filename );
}

sub loadVirtDB {
   return if $virtDB;
   $virtDB = &readDB( $domainDBName );
}

sub loadNodePopulations {
   return if $nodePopulations;
   $nodePopulations = &readDB( $nodePopulationDBName );
}

sub loadNodeDB {
   return if $nodeDB;
   $nodeDB = &readDB( $nodeDBName );
}

sub domain {
   my $action = lc shift;
   my $return = '';
   &loadVirtDB();
   &loadNodePopulations();
   @_ = keys( %$virtDB ) if ( $_[0] && $_[0] eq 'ALL' );
   if ( $_[0] && $_[0] eq 'RUNNING' ) {
      my @running;
      foreach my $node ( keys %$nodePopulations ) {
         push @running, keys %{ $nodePopulations->{$node}->{'running'} };
      }
      @_ = @running;
   }
   if ( $action eq 'update' ) { # download xml to var and update database
      while ( my $virt = shift ) {
         &parseDomain( $virt );
      } # while
      &writeDB( $domainDBName, $virtDB );
   } elsif ( $action eq 'list' ) { # dump domain as a tab separated data file
      my @return;
      foreach my $node ( keys %$nodePopulations ) {
         foreach my $virt (keys %{$nodePopulations->{$node}->{'running'}} ) {
            push @return, &listDomain( $virt, $node );
         }
      }
      $return = join( "\n", sort @return ) . "\n";;
   }
   return $return;;
} # sub domain

sub listDomain {
   my ($virt,$node) = @_;
   my @return;
   push @return, $virt;
   push @return, $node;
   foreach my $column ( sort keys %{ $virtDB->{$virt} } ) {
      push @return, $virtDB->{$virt}->{$column};
   }
   return join( "\t", @return);
}
   


# get the XML definition file of a running domain off of whatever
# node it is running on, and save it to disk
sub getVirtConfig {
   my ($virt,$filename) = @_;
   my $return;
   print "In getVirtConfig looking for $virt with file $filename\n" if $DEBUG;
   if ( -f $filename ) {
      open XML, "<$filename" or die "Could not read from $filename: $!\n";
      $return = join( '', <XML> );
      close XML;
   } else {
      &loadNodePopulations();
      #die Dumper( $nodePopulations );
      foreach my $node ( keys %$nodePopulations ) {
         print "getVirtConfig Looking on $node for $virt\n";
         if ( exists( $nodePopulations->{$node}->{'running'}->{$virt} ) ) { # we found it
            print "Found $virt on node $node\n";
            $return = `ssh $node 'virsh dumpxml $virt'`;
            open XML,">$filename" or die "Could not write to $filename: $!\n";
            print XML $return;
            close XML;
         } # if
      } # foreach
   } # if..else
   return $return;
} # sub getVirtConfig

sub getXMLValue {
   my ( $key, $string ) = @_;
   my $start = "<$key";
   my $end = "</$key>";
   $string =~ m/$start([^>]*)>([^<]+)$end/;
   return ($1,$2);
}

sub parseDomain {
   my ($virt, $nodePopulations ) = @_;
   
   my @keysToSave = ( 'uuid', 'memory', 'vcpu' );
   my $filename = "$confDir/$virt.xml";
   my $xml = &getVirtConfig( $virt, $filename );
   my ($param,$value) = &getXMLValue( 'uuid', $xml );
   $virtDB->{$virt}->{'uuid'} = $value;
   ($param,$value) = &getXMLValue( 'memory', $xml );
   $virtDB->{$virt}->{'memory'} = $value;
   ($param,$value) = &getXMLValue( 'vcpu', $xml );
   $virtDB->{$virt}->{'vcpu'} = $value;

   $xml =~ m/type='vnc' port='(\d+)'/;
   $virtDB->{$virt}->{'vnc'} = $1;
}

sub getDomainsOnNode {
   my $node = shift;
   my @nodeList = grep { /^\s*\d/ } `ssh $node 'virsh list'`;
   for ( my $i = 0; $i < @nodeList; $i++ ) {
      if ( $nodeList[$i] =~ m/\s*\d+\s*([^ ]+)/ ) {
         $nodeList[$i] = $1;
      }
   }
   my %hash = map{ $_ => time } @nodeList;
   return \%hash;
}

sub node {
   my $action = lc shift;

   my %conversion = ( 
     'CPU frequency' => 'clock',
     'CPU model' => 'cpu_model',
     'CPU socket(s)' => 'cpu_socket',
     'CPU(s)' => 'cpu_count',
     'Core(s) per socket' => 'cpu_cores',
     'Memory size' => 'memory',
     'NUMA cell(s)' => 'numa_cells',
     'Thread(s) per core' => 'threads_per_core'
   );


   print "In node, action is $action\n" if $DEBUG;
   my $return = '';
   &loadNodeDB();
   if ( $action eq 'update' ) { # read information for nodes and update database
      @_ = keys %$nodeDB if ( $_[0] eq 'ALL' );
      while ( my $nodename = shift ) {
         print "Updating $nodename\n" if $DEBUG;
         $return = `ssh $nodename 'virsh nodeinfo'`;
         print "Output of ssh $nodename 'virsh nodeinfo' is\n" . $return if $DEBUG;
         my @nodeinfo = split( "\n", $return );
         for ( my $i = 0; $i < @nodeinfo; $i++ ) {
            my ($key, $value) = split( /:\s+/, $nodeinfo[$i] );
            if ( $value =~ m/^(\d+)\s+[a-z]+$/i ) {
               $value = $1;
            }
            $key = $conversion{$key} if exists( $conversion{$key} );
            $nodeDB->{$nodename}->{$key} = $value;
         } # for
      } # while
      print "nodeDB state after update\n" . Dumper( $nodeDB ) if $DEBUG;
      &writeDB( $nodeDBName, $nodeDB );
   } elsif ( $action eq 'list' ) { # dump database as a tab separated file with headers
      my @return;
      foreach my $node ( sort keys %$nodeDB ) {
         @return[0] = "Node\t" . join( "\t", sort keys %{ $nodeDB->{$node} } ) unless @return;
         my @line;
         push @line, $node;
         foreach my $column (sort keys %{ $nodeDB->{$node} }) {
            push @line, $nodeDB->{$node}->{$column};
         }
         push @return, join( "\t", @line );
      }
      $return = join( "\n", @return ) . "\n";
   } elsif ( $action eq 'scan' ) {
      foreach my $node ( keys %$nodeDB ) {
         $nodePopulations->{$node}->{'running'} = &getDomainsOnNode( $node );
         $nodePopulations->{$node}->{'lastchecked'} = time;
      }
      &writeDB( $nodePopulationDBName,$nodePopulations );
   } # if..elsif
   return $return;
}


sub cluster {
   my $action = lc shift;
   my $return = '';
   if ( $action eq 'status' ) {
      &loadVirtDB();
      &loadNodePopulations();
      &loadNodeDB();
      print "Node\tThreads\tMemory\tDomains\tvcpu\tmem_used\n";
      my $usedmem = 0;
      my $usedcpu = 0;
      my $availmem = 0;
      my $availcpu = 0;
      my $totalDomains = 0;
      foreach my $node (sort keys %$nodeDB ) {
         my $memory = 0;
         my $vcpus = 0;
         my $count = 0;
         foreach my $domain ( keys %{ $nodePopulations->{$node}->{'running'} } ) {
            $memory += $virtDB->{$domain}->{'memory'};
            $vcpus += $virtDB->{$domain}->{'vcpu'};
            $count++;
         }
         $return .= "$node\t$nodeDB->{$node}->{cpu_count}\t$nodeDB->{$node}->{memory}\t$count\t$vcpus\t$memory\n";
         $usedmem += $memory;
         $usedcpu += $vcpus;
         $totalDomains += $count;
         $availmem += $nodeDB->{$node}->{memory};
         $availcpu += $nodeDB->{$node}->{cpu_count};
      } # outer for
      $return .= "Total\t$availcpu\t$availmem\t$totalDomains\t$usedcpu\t$usedmem\n";
   }
   return $return;
}

      
#my $config = &readConf( $confFile );

my $command = shift; # the first one is the actual subsection
my $action = shift; # second is action to run

if ( $command eq 'node' ) {
   print &node( $action, @ARGV );
} elsif ( $command eq 'domain' ) {
   print &domain( $action, @ARGV );
} elsif ( $command eq 'cluster' ) {
   print &cluster( $action, @ARGV );
} else {
   &help();
}


1;