#!/usr/bin/env perl
use warnings;
use strict;
use Net::OpenSSH;

my $home_dir = $ENV{HOME} or die "HOME environment variable not set\n";

# Per-hypervisor connect timeout in seconds. Override with TIMEOUT=10.
my $timeout = $ENV{TIMEOUT} // 60;

my %hypervisors = (
	"kellislab1.csail.mit.edu" => "kl1",
	"kellislab2.csail.mit.edu" => "kl2",
        "kellislab3.csail.mit.edu" => "kl3",
	"kellislab4.csail.mit.edu" => "kl4",
	"jibx" => "jx",
	"jiby" => "jy",
      );

my $vm_config  = "$home_dir/.config/ssh/vm.conf";
my $inv_dir    = "$home_dir/.local/share/dotfiles/ansible/inventory";
my $inv_file   = "$inv_dir/vms.yml";

open my $fh, '>', $vm_config or die "Cannot create SSH config at $vm_config: $!\n";
print $fh "# Auto-generated VM config - ", scalar(localtime), "\n\n";

# Per-hypervisor group -> [VM aliases]. Collected during the DHCP sweep
# and flushed to an Ansible inventory file at the end. 
my %inv_groups;

sub write_inventory {
	my @hvs = sort keys %inv_groups;
	my $any = @hvs ? 1 : 0;
	unless (-d $inv_dir) {
		require File::Path;
		File::Path::make_path($inv_dir);
	}
	open my $inv, '>', $inv_file or do {
		warn "Cannot create inventory at $inv_file: $!\n";
		return;
	};
	print $inv "# Auto-generated by update_vm_openssh - ", scalar(localtime), "\n";
	print $inv "# DO NOT EDIT. Re-run update_vm_openssh to refresh.\n";
	print $inv "# Static hosts/groups belong in hosts.yml (same directory).\n\n";
	unless ($any) {
		print $inv "# No VMs discovered on any hypervisor.\n";
		close $inv;
		return;
	}
	print $inv "all:\n";
	print $inv "  children:\n";
	print $inv "    vms:\n";
	print $inv "      children:\n";
	for my $hv (@hvs) {
		print $inv "        $hv:\n";
		print $inv "          hosts:\n";
		for my $alias (sort @{$inv_groups{$hv}}) {
			print $inv "            $alias:\n";
		}
	}
	close $inv;
}

for my $hv (sort keys %hypervisors) {
	my $ssh = Net::OpenSSH->new($hv, timeout => $timeout) or do {
		warn "Failed to connect to $hv\n";
		next;
	};
	# Fetch DHCP leases for default network
	my @lease_lines = $ssh->capture(
		{stdin_data => ''},
		'LIBVIRT_DEFAULT_URI=qemu:///system virsh net-dhcp-leases default');

	if ($ssh->error) {
		warn "virsh net-dhcp-leases failed on $hv\n";
		next;
	}

	# Fetch defined domains so stale DHCP leases (from destroyed VMs whose
	# leases haven't expired in dnsmasq yet) are filtered out. 
	my @domain_lines = $ssh->capture(
		{stdin_data => ''},
		'LIBVIRT_DEFAULT_URI=qemu:///system virsh list --all --name');
	if ($ssh->error) {
		warn "virsh list --all failed on $hv; skipping stale-lease filter\n";
		@domain_lines = ();
	}
	my %defined_domains = map { chomp; $_ => 1 } grep { /\S/ } @domain_lines;
	my $header_line;
	for my $line (@lease_lines) {
		if ($line =~ /^\s*Expiry/i &&
		    $line =~ /\bIP address\b/ &&
		    $line =~ /\bHostname\b/) {
			$header_line = $line;
			last;
		}
	}
	unless ($header_line) {
		warn "[$hv] Could not locate header with IP address & Hostname. Skipping.\n";
		next;
	}
	my $ip_pos   = index($header_line, 'IP address');
	my $host_pos = index($header_line, 'Hostname');
	my $client_pos = index($header_line, 'Client ID');

	if ($ip_pos < 0 || $host_pos <= $ip_pos) {
		warn "[$hv] Malformed header: IP/Hostname positions invalid. Skipping.\n";
		next;
	}
	# Calculate column widths
	my $ip_len   = $host_pos - $ip_pos;
	my $host_len = ($client_pos > $host_pos) 
	  ? $client_pos - $host_pos 
	  : (length($header_line) - $host_pos); # Fallback to end of line
	my %lease_map;
	for my $line (@lease_lines) {
		next if $line =~ /^\s*$/ ||
		  $line =~ /^-+/ ||
		  index($line, 'IP address') == $ip_pos;
		# Extract fixed-width columns based on header positions
		my $ip_raw   = substr($line, $ip_pos, $ip_len);
		my $host_raw = substr($line, $host_pos, $host_len);

		# Trim whitespace
		$ip_raw   =~ s/^\s+|\s+$//g;
		$host_raw =~ s/^\s+|\s+$//g;

		# Skip invalid/placeholder entries
		next unless $host_raw && $host_raw ne '-' && $ip_raw;

		# Skip stale leases: the VM was destroyed but dnsmasq hasn't
		# expired the entry yet. Cross-reference against defined domains.
		if (%defined_domains && !$defined_domains{$host_raw}) {
			warn "[$hv] Skipping stale lease for '$host_raw' (domain not defined)\n";
			next;
		}

		# Strip CIDR notation (e.g., /24)
		$ip_raw =~ s{[/]\d+$}{};
		$lease_map{$host_raw} = $ip_raw;

	}
	for my $vm (sort keys %lease_map) {
		my $alias = "$vm.$hypervisors{$hv}";
		printf $fh "Host %s.%s %s\n",   $vm, $hv, $alias;
		printf $fh "    HostName %s\n", $lease_map{$vm};
		printf $fh "    ProxyJump %s\n",  $hv;
		printf $fh "    ControlMaster auto\n";
		printf $fh "    ControlPersist 1800\n";
		print  $fh "    ControlPath ~/.ssh/cm_socket/%C\n";
		print  $fh "    User agent\n\n";
		push @{$inv_groups{$hypervisors{$hv}}}, $alias;
	} 
}
close $fh;
print "SSH config successfully updated at $vm_config\n";

write_inventory();
if (-e $inv_file) {
	print "Ansible inventory written to $inv_file\n";
}
