#!/usr/bin/perl

use strict;
use warnings;

use AptPkg::Cache;
use Dpkg::Version;
use Getopt::Long;
use List::Util qw(uniqstr);

use PVE::Tools qw(run_command);
use PVE::SysFSTools;
use PVE::SafeSyslog;

## Module-Wide Variables & Constants

my @BASE_PACKAGE_DEPENDENCIES = qw(dkms gcc libc6-dev proxmox-default-headers);

my $apt_cache; # initialized in the main section

my ($no_block_nouveau, $help); # CLI Options, initialized in the main section

## Helpers

my sub print_usage_and_exit {
    my ($error_message) = @_;
    print("$error_message\n\n") if $error_message;
    print(<<"EOF"
USAGE:   pve-nvidia-vgpu-helper [OPTIONS]
 Commands:
    setup                 Install dependencies and ensure the node is ready for using NVIDIA vGPU.
      --no-block-nouveau  Do not add modprobe.d config entry to block loading the 'nouveau' module.

 General Options:
    --help                This output.
EOF
    );
    exit($error_message ? 1 : 0);
}

my sub package_is_installed {
    my ($package) = @_;
    my $p = $apt_cache->{$package};
    return defined($p->{CurrentState}) && $p->{CurrentState} eq "Installed";
}

my sub get_running_proxmox_kernel_version {
    my $running_kernel_raw;
    run_command(['/usr/bin/uname', '-r' ], outfunc => sub { $running_kernel_raw = shift } );

    if ($running_kernel_raw =~ m/^(\d+\.\d+\.\d+-\d+)-pve$/) {
	return $1;
    }
    return;
}

my sub apt_install {
    my (@packages) = @_;

    if (system('apt-get', '--no-install-recommends', 'install', '--', @packages) != 0) {
	die "apt failed during the installation (exit code $?), please check its output for errors.\n";
    }
}

## Commands

my sub do_setup {
    if ($no_block_nouveau) {
	print "Skipping creation of modprobe.d config to block 'nouveau' module as requested\n";
    } elsif (! -e "/etc/modprobe.d/block-nouveau.conf") {
	print "Adding modprobe config to block loading of the 'nouveau' module.\n";
	mkdir('/etc/modprobe.d') or $!{EEXIST} or die "failed to create '/etc/modprobe.d' - $!\n";
	PVE::SysFSTools::file_write("/etc/modprobe.d/block-nouveau.conf", "blacklist nouveau\n")
	  || die "failed to create block-nouveau.conf - $!\n";
	syslog('info', "Added modprobe config to block loading of the 'nouveau' module.");

	print "Updating initramfs to deploy modprobe.d change...\n";
	run_command(["update-initramfs", "-u", "-k", "all"]);
    } else {
	print "No need to create 'block-nouveau.conf' modprobe.d config as it already exists.\n";
    }

    my @missing_packages = grep { !package_is_installed($_) } @BASE_PACKAGE_DEPENDENCIES;

    my $running_proxmox_kernel_version = get_running_proxmox_kernel_version();
    if (defined($running_proxmox_kernel_version)) {
	print "You are running the Proxmox kernel $running_proxmox_kernel_version, searching"
	    ." the associated and newer kernel headers package.\n";

	run_command(
	    ['/usr/bin/dpkg-query', '-W', '-f', '${db:Status-Status}\t${binary:Package}\n', 'proxmox-kernel-*-pve*'],
	    outfunc => sub {
		my $kernel = shift;

		return if $kernel !~ m/^installed\s+proxmox-kernel-((\d+.\d+).\d+-\d+)-pve(?:-signed)?\s*$/;

		my ($full_version, $major_minor) = ($1, $2);
		if (Dpkg::Version::version_compare($running_proxmox_kernel_version, $full_version) != 1) {
		    my $header_package_name = "proxmox-headers-$major_minor";
		    if (!package_is_installed($header_package_name)) {
			push @missing_packages, $header_package_name;
		    }
		}
	    },
	);
    } else {
	print "You are not currently running a proxmox kernel, so we cannot determine the"
	    ." appropriate kernel header package.\n"
	    ." Please make sure you have the appropriate header package installed.\n";
    }

    if (@missing_packages){
	@missing_packages = sort(uniqstr(@missing_packages)); # deduplicate
	print "The following packages are missing:\n" . join("\n", map { "\t$_"} @missing_packages) ."\n";
	print "Would you like to install them now (y/N)? ";

	my $answer = <STDIN>;
	if (defined($answer) && $answer =~ m/^\s*y(?:es)?\s*$/i) {
	    apt_install(@missing_packages);
	} else {
	    print "Aborted by user request\n";
	    return;
	}
    } else {
	print "All required packages are already installed.\n"
    }

    print "All done, you can continue with the NVIDIA vGPU driver installation.\n";
}

## Main

GetOptions(
    'no-block-nouveau' => \$no_block_nouveau,
    'help' => \$help,
) or print_usage_and_exit("Could not parse commandline options");

my $command = shift // print_usage_and_exit('Missing command.');
if ($help) {
    print_usage_and_exit();
}

die "Please execute the script with root privileges\n" if $> != 0;

$apt_cache = AptPkg::Cache->new();
die "unable to initialize AptPkg::Cache\n" if !$apt_cache;

if ($command eq 'setup') {
    do_setup();
} else {
    print_usage_and_exit("unknown command '$command'");
}
