#!/usr/bin/perl -w

use constant {
    TIMEZONE => 'Europe/Amsterdam'
};

use strict;
use Getopt::Long;
use PVE::APIClient::LWP;
use PVE::AccessControl;
use PVE::INotify;
use JSON;
use DateTime;
use Data::Dumper;

sub usage {
    my $error = shift;

    if (defined($error)) {
        print "ERROR: $error\n\n\n";
    }

    print "
This script will balance VM's over your cluster. It will look at memoryusage on
all nodes of your cluster, and move memory hungry VM's to a less busy node. 

Options:
  pmbalance [--dry-run] [--verbose] [--email]
    --dry-run  Don't actually move nodes, show what you would move
    --verbose  Add some extra output
    --email    Send an email to <emailaddress> for move
    --help     This help

";

    if(defined($error)) {
        exit(1);
    }
}

my (%todo, $dryrun, $verbose, $email, $help);
my $timelimit = DateTime->now(time_zone => TIMEZONE);

GetOptions (
    "dry-run"    => \$dryrun,
    "verbose"    => \$verbose,
    "help"       => \$help,
    "email=s"    => \$email ) or die(usage());

if (defined($help)) {
    usage();
    exit(0);
}

my $hostname = PVE::INotify::read_file("hostname");

# normally you use username/password,
# but we can simply create a ticket and CRSF token if we are root
my $ticket = PVE::AccessControl::assemble_ticket('root@pam');
my $csrftoken = PVE::AccessControl::assemble_csrf_prevention_token('root@pam');

my $conn = PVE::APIClient::LWP->new(
    ticket => $ticket,
    csrftoken => $csrftoken);

sub verbose($) {
    my $msg = shift;

    if (!defined($verbose)) {
        return;
    }

    print "V: $msg\n";
}

sub get_nodes() {
    my %params;
    verbose("Looking for nodes on this cluster..");

    $params{'type'} = 'node';
    my $res = $conn->get("/api2/json/cluster/resources", \%params);

    my %nodes;
    foreach my $node (@{$res}) {
        my $name = $node->{'node'};
        verbose("Found $name");
        $nodes{$name}{'name'} = $name;
        $nodes{$name}{'free'} = $node->{'maxmem'}-$node->{'mem'};
        $nodes{$name}{'max'}  = $node->{'maxmem'};
        $nodes{$name}{'used'} = $node->{'mem'};
        $nodes{$name}{'percfree'} = int($nodes{$name}{'free'}/($node->{'maxmem'}/100));
    }

    return \%nodes;
}

sub get_vm_description($$) {
    my $node = shift;
    my $vmid = shift;

    my %params;
    $params{'current'} = '1';
    my $res = $conn->get("/api2/json/nodes/$node/qemu/$vmid/config", \%params);

    if (defined($res->{'description'})) {
        return $res->{'description'};
    }

    return undef;
}

sub get_vm_to_move($$) {
    my $sender = shift;
    my $memsize = shift;

    my %vms;

    verbose("Looking for migratable VMs on $sender");
    my %params;
    $params{'type'} = 'vm';
    my $res = $conn->get("/api2/json/cluster/resources", \%params);
    foreach my $vm (@{$res}) {
        my $vmid = $vm->{'vmid'};
        # VM Must be on sender node
        if ($vm->{'node'} ne $sender) { next; };
        # VM Must be running
        if ($vm->{'status'} ne "running") { next; };
        # VM Must not be too large
        if ($vm->{'maxmem'} > $memsize) { next; };
        # VM Must not be pinned to this host
        my $description = get_vm_description($sender, $vmid);
        if (defined($description)) {
            if ($description =~ m/Pinned: $sender/g) {
                next;
            }
        }
        $vms{$vmid} = $vm->{'maxmem'};
    }

    my @keys = sort { $vms{$b} <=> $vms{$a} } keys %vms;
    verbose("Returning the largest vm: ".$keys[0]);
    return $keys[0];
}

sub get_avg_free($) {
    my $nodes = shift;

    verbose("Calculating the average memory free on the cluster");
    my $total = 0;
    for my $node (keys(%$nodes)) {
        verbose("$node is at ".$nodes->{$node}{'percfree'}."% free");
        $total += $nodes->{$node}{'percfree'};
    }

    verbose("Which is: ".int($total/keys(%$nodes))."%");

    return int($total/keys(%$nodes));
}

sub get_donor_recipient() {
    my $nodes = get_nodes();
    my $avgmem = get_avg_free($nodes);
    my $sender = {
        "name"  => undef,
        "percfree" => 100 };

    my $receiver = {
        "name" => undef,
        "percfree" => 0 };

    for my $node (keys(%$nodes)) {
        if ($nodes->{$node}{'percfree'} < $avgmem and $avgmem-$nodes->{$node}{'percfree'} >= 5) {
            if ($nodes->{$node}{'percfree'} < $sender->{'percfree'}) {
                $sender = $nodes->{$node};
            }
        } else {
            if ($nodes->{$node}{'percfree'} > $receiver->{'percfree'}) {
                $receiver = $nodes->{$node};
            }
        }
    }

    if (!defined($sender->{'name'})) {
        print "Nothing to do. No nodes much more busy than average...\n";
        return {
            'sender' => undef};
    }

    my $diff = $avgmem-$sender->{'percfree'};
    my $bdiff = int(($sender->{'max'}/100)*$diff);
    verbose("We should lower usage of ".$sender->{'name'}." by: ".$diff."% to reach ".$avgmem."%");
    verbose($diff." of ".$sender->{'max'}." = ".$bdiff." bytes");
    verbose($receiver->{'name'}." will receive the VM");

    return {
        'sender' => $sender->{'name'},
        'receiver' => $receiver->{'name'},
        'bdiff' => $bdiff }
}


my $lbstate = get_donor_recipient();
if (!defined($lbstate->{'sender'})) {
    exit(0);
}


my $tomove = get_vm_to_move($lbstate->{'sender'}, $lbstate->{'bdiff'});

my %params;
$params{'target'} = $lbstate->{'receiver'};
$params{'online'} = 1;

if (!defined($dryrun)) {
    print "Move $tomove from ".$lbstate->{'sender'}." to ".$lbstate->{'receiver'}."\n";
    my $res = $conn->post("/api2/json/nodes/".$lbstate->{'sender'}."/qemu/".$tomove."/migrate", \%params);
} else {
    print "We would move $tomove from ".$lbstate->{'sender'}." to ".$lbstate->{'receiver'}."\n";
}