#!/usr/bin/env perl

use v5.12;
use strict;
use warnings;

use Config::Tiny;
use File::Basename;
use File::Spec::Functions qw(catdir catfile);
use Getopt::Long qw(:config no_ignore_case gnu_getopt auto_help);
use List::Util qw(first);
use Switch;

use constant {
  VERSION => "0.3.1"
};

sub unquote($) {
  return ($_[0] =~ m/"?([^"]*)"?/)[0];
}

my $quiet = '';
my $output_file = '';
my $version = 0;
GetOptions('output-file|output|o=s' => \$output_file,
           'version' => \$version,
           'quiet|q' => \$quiet)
  or die "invalid options";

if ($version) {
  say 'snapman-grub version '.VERSION;
  exit
}


my $_grubconf = Config::Tiny->read('/etc/default/grub')
  or die "could not load grub default config";
my %grubconf = %{$_grubconf->{_}};

my $_snapconf = Config::Tiny->read('/etc/snapman.d/snapman.conf')
  or die "could not load snapman config file";
my %snapconf = %{$_snapconf->{grub}};

my %grub_info_cache = ();

my $snapman = 'snapman';
my $distributor = unquote($grubconf{GRUB_DISTRIBUTOR});
my @kernel_order = split(/\s/, unquote($snapconf{kernel_order} // ''));
my @snapshot_order = split(/\s/, unquote($snapconf{snapshot_order} // ''));
my $other_kernels_switch = unquote($snapconf{other_kernels} // 'submenu');
my $other_snapshots_switch = unquote($snapconf{other_snapshots} // 'submenu');
my $date_snapshots_switch = unquote($snapconf{date_snapshots} // 'submenu');

# check if variables have sane values
my @yesNoSubmenu = ('yes', 'no', 'submenu');
die "configuration variable 'other_kernels' has invalid value.\nwas '${other_kernels_switch}', must be one of ('yes', 'no', 'submenu')"
  unless grep {$other_kernels_switch eq $_} @yesNoSubmenu;
die "configuration variable 'other_snapshots' has invalid value.\nwas '${other_snapshots_switch}', must be one of ('yes', 'no', 'submenu')"
  unless grep {$other_snapshots_switch eq $_} @yesNoSubmenu;
die "configuration variable 'date_snapshots' has invalid value.\nwas '${date_snapshots_switch}', must be one of ('yes', 'no', 'submenu')"
  unless grep {$date_snapshots_switch eq $_} @yesNoSubmenu;

my $indentation = 0;


sub grubDevice($) {
  my $ret = `grub-probe -t device $_[0]`;
  chomp $ret;
  return $ret;
}

sub grubRelPath($) {
  my $ret = `grub-mkrelpath $_[0]`;
  chomp $ret;
  return $ret;
}

sub grubHints($) {
  my $ret = `grub-probe -t hints_string $_[0]`;
  chomp $ret;
  return $ret;
}

sub grubPartmap($) {
  my $ret = `grub-probe -t partmap $_[0]`;
  chomp $ret;
  return $ret;
}

sub grubFsUuid($) {
  my $ret = `grub-probe -t fs_uuid $_[0]`;
  chomp $ret;
  return $ret;
}

sub grubCompatibilityHint($) {
  my $ret = `grub-probe -t compatibility_hint $_[0]`;
  chomp $ret;
  return $ret;
}

sub printIndented(@) {
  print "\t" x $indentation;
  print @_;
}

sub sayIndented(@) {
  printIndented @_;
  say "";
}

sub printPartmapModule(+) {
  my ($snap) = @_;
  my ($root_snap, $boot_snap) = @{$snap};

  # translated from grub-mkconfig_lib
  foreach my $partmodule (@{$grub_info_cache{$boot_snap ? 'boot_partmap' : 'root_partmap'}}) {
    switch ($partmodule) {
      case ["netbsd", "openbsd"] { sayIndented "insmod part_bsd" }
      else { sayIndented "insmod part_${partmodule}" }
    }
  }
}


sub printPlatformHints(+$) {
  my ($snap, $kernel_path) = @_;
  my ($root_snap, $boot_snap) = @{$snap};
  my $fs_uuid = $grub_info_cache{$boot_snap ? 'boot_fs_uuid' : 'root_fs_uuid'};
  my $hints = $grub_info_cache{$boot_snap ? 'boot_search_hints' : 'root_search_hints'};
  my $template = <<END;
if [ x\$feature_platform_search_hint = xy ]; then
  search --no-floppy --fs-uuid --set=root ${hints} ${fs_uuid}
else
  search --no-floppy --fs-uuid --set=root ${fs_uuid}
fi
END

  for my $line ($template =~ m/^(.*)\n/mg) {
    sayIndented $line;
  }
}


sub printLinuxLine($$) {
  my ($kernel_path, $linux_cmdline) = @_;
  my $kernel_path_rel = grubRelPath $kernel_path;
  printIndented "linux ${kernel_path_rel}";

  print " root=UUID=".($grub_info_cache{root_fs_uuid});
  print " rw ".$linux_cmdline;
  say "";
}


sub linuxEntry($$$$$) {
  my ($snap, $kernel_path, $kernel_base_name, $initrd_path, $initrd_type) = @_;
  my ($root_snap, $boot_snap) = @{$snap};
  my $root_snap_relpath = grubRelPath $root_snap;

  printIndented ("menuentry '${distributor} Linux, snapshot ${root_snap_relpath} with Linux ${kernel_base_name}");

  switch ($initrd_type) {
    case "fallback" { print " (fallback image)"; }
    case "recovery" { print " (recovery mode)"; }
  }


  print "' --class gnu-linux --class gnu --class os";
  say " {";
  $indentation++;

  sayIndented "load_video";
  sayIndented "set gfxpayload=$grubconf{GRUB_GFXPAYLOAD_LINUX}"
    unless !defined($grubconf{GRUB_GFXPAYLOAD_LINUX});

  sayIndented "insmod gzio";
  sayIndented "insmod btrfs"; # doesn't make sense with another file system

  printPartmapModule $snap;

  #my $compathint = $boot_snap ? $grub_info_cache->{boot_compathint} : $grub_info_cache->{root_compathint};
  my $compathint = $grub_info_cache{$boot_snap ? 'boot_compathint' : 'root_compathint'};
  sayIndented ("set root=".$compathint)
    if $compathint;

  my $linux_cmdline = (unquote $grubconf{GRUB_CMDLINE_LINUX})." rootflags=subvol=".grubRelPath $root_snap;
  $linux_cmdline = $linux_cmdline." ".(unquote $grubconf{GRUB_CMDLINE_LINUX_DEFAULT})
    unless $initrd_type eq "recovery";

  printPlatformHints($snap, $kernel_path);

  sayIndented "echo 'Loading Linux ${kernel_base_name} ...'";
  printLinuxLine($kernel_path, $linux_cmdline);

  sayIndented "echo 'Loading initial ramdisk ...'";
  my $initrd_path_rel = grubRelPath $initrd_path;
  my $initrd_str = $initrd_path_rel;
  my $bootdirname = dirname $initrd_path;#$initrd_path =~ m|^(([^/]*/)*)|s;
  my $bootdirname_rel = dirname $initrd_path_rel;#$initrd_path =~ m|^(([^/]*/)*)|s;
  $initrd_str = catfile($bootdirname_rel, "intel-ucode.img")." ".$initrd_path_rel
    if -e catfile($bootdirname, "intel-ucode.img");
  sayIndented "initrd ${initrd_str}";

  $indentation--;
  print "\t" x $indentation;
  say "}\n";
}

sub rateInitrdName($) {
  my ($name) = @_;

  return 2 if ($name =~ /recovery\.img/);
  return 1 if ($name =~ /fallback\.img/);
  return 0;
}

sub kernelFileName($) {
  my ($kernel_path) = @_;
  return ($kernel_path =~ m|/([^/]*)$|)[0];
}

sub kernelBaseName($) {
  my ($kernel) = @_;
  return ($kernel =~ m|^[^-]*-(.*)$|)[0];
}

sub findAndRemove(&\@) {
  my ($fun,$list) = @_;

  for (my $i = 0; $i < @{$list}; $i++) {
    my $curr = $list->[$i];
    if (&$fun($curr)) {
      splice(@{$list}, $i, 1);
      return $curr;
    }
  }

  return undef;
}

sub handleKernel($$$) {
  my ($snap, $kernel_path, $grub_info_cache) = @_;
  say STDERR "found kernel: ${kernel_path}";
  my $kernel_dir = dirname $kernel_path;
  my $kernel = kernelFileName $kernel_path;
  my $kernel_base_name = kernelBaseName $kernel;

  my @potential_initrds = sort {(rateInitrdName $a) <=> (rateInitrdName $b)} <${kernel_dir}/init*>;
  for my $initrd_path (@potential_initrds) {
    next unless $initrd_path =~ /${kernel_base_name}(-fallback|-recovery|)*\.img/;
    my ($initrd_type) = $initrd_path =~ m/${kernel_base_name}-*((fallback|recovery|)*)\.img/;

    next if ($initrd_type eq "-recovery") && $grubconf{GRUB_DISABLE_RECOVERY};

    say STDERR "found initrd: ${initrd_path}";

    linuxEntry($snap, $kernel_path, $kernel_base_name, $initrd_path, $initrd_type);
  }
}


# snap is an array ref. ($root_snap, $boot_snap)
sub handleSnapshot($\%) {
  my ($snap, $grub_info_cache) = @_;
  my ($root_snap, $boot_snap) = @{$snap};
  say STDERR "subvolume found: ".$root_snap;

  my $kernel_dir = $boot_snap || catdir($root_snap, 'boot');

  my @other_kernels = <${kernel_dir}/vmlinu*>;
  my @ordered_kernels = ();

  # the user wanted some kernels above others
  foreach my $tmp (@kernel_order) {
    my $regex = qr/$tmp$/;
    my $krnl = findAndRemove {$_[0] =~ /$regex/} @other_kernels;
    push(@ordered_kernels, $krnl)
      if defined $krnl;
  }

  foreach my $kernel_path (@ordered_kernels) {
    handleKernel($snap, $kernel_path, $grub_info_cache);
  }

  if (($other_kernels_switch ne "no") && (@other_kernels > 0)) {
    if ($other_kernels_switch eq "submenu") {
      sayIndented "submenu 'Other kernels for ${distributor} Linux, snapshot $root_snap' {";
      $indentation++;
    }

    foreach my $kernel_path (@other_kernels) {
      handleKernel($snap, $kernel_path, $grub_info_cache);
    }

    if ($other_kernels_switch eq "submenu") {
      $indentation--;
      sayIndented "}";
    }
  }
}

if ($output_file ne '') {
  open(STDOUT, '>', $output_file)
    or die $!;
}

if ($quiet) {
  open(STDERR, '>', '/dev/null')
    or die $!;
}

my @other_snaps = map {[split /\t/, $_]} (`$snapman list --fullpath --named` =~ m/(^.*)$/mg);
my @date_snaps = map {[split /\t/, $_]} sort {$b cmp $a} (`$snapman list --fullpath --time` =~ m/(^.*)$/mg);
my @ordered_snaps = ();
my $snapshot_names = join(' ', @snapshot_order);
my @snapshot_order_fullname = (`$snapman printname -B "" --fullpath ${snapshot_names}` =~ m/^(.*)$/mg);

my $first_root_snapshot = $other_snaps[0]->[0] || $date_snaps[0]->[0];
my @bootsnap_list = map {$_->[1]} @other_snaps;
push(@bootsnap_list, map {$_->[1]} @date_snaps);
my $first_boot_snapshot = first {$_} @bootsnap_list;

%grub_info_cache = (
  root_fs_uuid => grubFsUuid($first_root_snapshot),
  boot_fs_uuid => $first_boot_snapshot ? grubFsUuid($first_boot_snapshot) : '',
  root_partmap => [split / /, grubPartmap($first_root_snapshot)],
  boot_partmap => $first_boot_snapshot ? [split / /, grubPartmap($first_boot_snapshot)] : [],
  root_compathint => grubCompatibilityHint($first_root_snapshot),
  boot_compathint => $first_boot_snapshot ? grubCompatibilityHint($first_boot_snapshot) : '',
  root_search_hints => grubHints($first_root_snapshot),
  boot_search_hints => $first_boot_snapshot ? grubHints($first_boot_snapshot) : '');


for my $tmp (@snapshot_order_fullname) {
  my $regex = qr/^${tmp}$/;
  my $snpsht = findAndRemove {$_[0]->[0] =~ /$regex/} @other_snaps;
  push(@ordered_snaps, $snpsht)
    if defined $snpsht;
}

foreach my $snap (@ordered_snaps) {
  handleSnapshot($snap, %grub_info_cache);
}

if ($other_snapshots_switch ne "no" && @other_snaps > 0) {
  if ($other_snapshots_switch eq "submenu") {
    sayIndented "submenu 'Other snapshots for ${distributor} Linux' {";
    $indentation++;
  }

  foreach my $snap (@other_snaps) {
    handleSnapshot($snap, %grub_info_cache);
  }

  if ($other_snapshots_switch eq "submenu") {
    $indentation--;
    sayIndented "}";
  }
}

if ($date_snapshots_switch ne "no" && @date_snaps > 0) {
  if ($date_snapshots_switch eq "submenu") {
    sayIndented "submenu 'Date-snapshots for ${distributor} Linux' {";
    $indentation++;
  }

  foreach my $snap (@date_snaps) {
    handleSnapshot($snap, %grub_info_cache);
  }

  if ($date_snapshots_switch eq "submenu") {
    $indentation--;
    sayIndented "}";
  }
}

if ($output_file ne '') {
  close STDOUT
    or die $!;
}
