#!/usr/bin/perl -w
use strict;
use vars qw($opt_q $opt_k $opt_f);
use Getopt::Std;

# Script to create a X11 Virtual frame buffer for the user
my $quiet = 0;
my $force_new = 0;
getopts('qk:f');
$quiet = 1 if (defined $opt_q);
$force_new = 1 if (defined $opt_f);

my $HOSTTYPE = `uname -s`;
chomp $HOSTTYPE;

my ($XVFB, $PS, $X11, $XVNC);
if ( $HOSTTYPE eq "Linux" ) {
    $XVFB='/usr/bin/Xvfb';
    $PS='/bin/ps';
    $X11='/usr/bin/Xorg';
    $XVNC='Xvnc';
} elsif ( $HOSTTYPE eq "Darwin" ) {
    $XVFB='/usr/X11R6/bin/Xvfb';
    $PS='/bin/ps';
    $X11='/Applications/Utilities/X11.app/Contents/MacOS/X11';
    $XVNC='Xvnc';
} else {
    print "Unsupported platform ($HOSTTYPE)\n";
    exit 1;
}

my $ME=$<;

# Find out if the user has a Xvfb server running
my @XVFBs_running=`$PS -Awwo uid,pid,command | grep Xvfb | grep -v grep`;

my (%displays, %users, %pids);

foreach my $displ_number (@XVFBs_running) {
    chomp $displ_number;
    $displ_number =~ s/^\s+//;
    $displ_number =~ s/\s+$//;
    $displ_number =~ s/\s{2,}/ /;

    my ($uid, $pid, $null, $display) = split / /, $displ_number;
    $display =~ s/:(\d+)/$1/;

    $displays{$display} = $uid;
    $users{$uid} += $display;
    $pids{$display} = $pid;
}
# Locate any running X11 displays
my @X11s_running=`$PS -Awwo uid,pid,command | grep $X11 | grep -v grep`;

foreach my $displ_number (@X11s_running) {
    chomp $displ_number;
    $displ_number =~ s/^\s+//;
    $displ_number =~ s/\s+$//;
    $displ_number =~ s/\s{2,}/ /;
    my ($uid, $pid, $null, $display, $null_auth, $null_psn);
    if ( $HOSTTYPE eq "Linux" ) {
	($uid, $pid, $null, $display) = split / /, $displ_number;
	
    } elsif ($HOSTTYPE eq "Darwin" ) {
	($uid, $pid, $null, $null_auth, $null_psn, $display) = split / /, $displ_number;
    }
    $display =~ s/:(\d+)/$1/;
    $displays{$display} = $uid;
    print "Cannot start - X11 display running on $display\n";
    exit 1;
}
# Locate any running VNC displays
my @XVNCs_running=`$PS -Awwo uid,pid,command | grep $XVNC | grep -v grep`;

foreach my $displ_number (@XVNCs_running) {
    chomp $displ_number;
    $displ_number =~ s/^\s+//;
    $displ_number =~ s/\s+$//;
    $displ_number =~ s/\s{2,}/ /;
    my ($uid, $pid, $null, $display) = split / /, $displ_number;
    $display =~ s/:(\d+)/$1/;
    $displays{$display} = $uid;
}

if ( $opt_k ) {
    # Kill a running server - $opt_k contains the display number
    if ( ! defined $displays{$opt_k} || ! defined $pids{$opt_k} ) {
	print "No display found on port $opt_k\n";
	exit 1;
    } elsif ( $displays{$opt_k} == $ME && defined $pids{$opt_k} ) {
	# We own the display
	print "Killing $opt_k...\n";
	kill 'TERM', $pids{$opt_k};
    } else {
	print "Display $opt_k not owned by you!\n";
	exit 1;
    }
} else {
    if ( defined $users{$ME} && ! $force_new) {
	# We have a server running, tell the user which one
	print "You already have the following displays open:\n" unless $quiet;
    
	foreach my $display (keys %displays) {
	    print "$display\n" if ( $displays{$display} == $ME );
	}
    } else {
	# No server running, so start one
	# Find a display number
	my $display = 1;
	while ( $displays{$display} ) {
	    $display ++;
	}
	my $failure=system "$XVFB :$display 2>&1 > /dev/null &";
	if ( ! $failure ) {
	    if ( $quiet ) {
		print "$display\n";
	    } else {
		print "Xvfb display started on :$display\n";
	    }
	} else {
	    print "$!\n ($failure)";
	    exit $failure;
	}
    }
    exit 0;
}
