#!/usr/local/bin/perl
#
#Add user to an existing project

use strict;
use warnings;
use lib '/cell_root/software/HPCC/lib/site_perl';
#For testing, use RW replica
#use lib '/afs/.glue.umd.edu/software/HPCC/1.1/lib/site_perl';

use Glue::HPCC::Cluster;
use Glue::HPCC::ClusterUtil;
use Glue::HPCC::Project::Person;
use Glue::HPCC::Project::Project;
use Glue::HPCC::Project::PersonProject;
use Glue::HPCC::Project::ProjectType;
use Glue::HPCC::Project::Log;
use Glue::HPCC::Utils::UserAccounts qw(user_exists);
use Glue::HPCC::Utils::MailUtils qw(send_template_email);

use Getopt::Long;

my $DEFAULT_CLOBBER_DEFACCT = 0;

my @cc_on_everything = ( 'payerle@umd.edu' );

sub usage_short()
{	print STDERR <<EOF;
$0: Add one or more users to one or more projects in cluster

Usage:
$0 [ -cluster CLUSTER ] -project PROJNAME -user USER -reason REASON  \
	[ -dryrun ] [ -verbose ] [ -[no]email ] [ -[no]highpriority ] \
	[ --default-account [high|standard|none] ] [ -[no]clobber-default-account ] \
	[ --FIXME ]

Adds one or more users to one or more pre-existing projects in specified
cluster.  You can repeat project, and user flags to allow
multiple projects, and/or users.

Or give 
$0 --help
for more help.
EOF
}

sub usage_long()
{	usage_short;

	my $tmp_clobber = 'noclobber-default-account';
	$tmp_clobber = 'clobber-default-account' if $DEFAULT_CLOBBER_DEFACCT;

	print STDERR <<EOF

Cluster will default based on where you are running the code.  Unless
the user already has home directory, etc. on the requested cluster, you
MUST run this command as root on the cluster headnode.

If the user does not already exist in the PDB person table, it will be
added with email of username\@umd.edu and fullname, uid from hesiod.  If
you want to change this, either edit the user afterwards or add to person
table manually.

This script will add to person table, add to appropriate netgroups, create
the user's home directory, and grant access to the appropriate accounts for
this project, all as needed.

Reason is the reason the project being added; typically an ITSC number or
reference to a AAC request.  It is required. 

The parameter --default-account determines what to set the user's default
account to.  If set to 'no' or 'none', the user's default account will NOT
be explicitly set; for an existing user this means it will remain unchanged,
for a new user Slurm will choose a default (probably std priority?).  If
'hi' or 'high', we will set the default account for the user to the high
priority account for this project (and an error will arise if no such account
exists).  If 'std' or 'standard', it will be set to the standard priority
account for this project.  If it is omitted, it will default to 'high' if
the project has a high priority account, and otherwise default to 'standard'.
Whether specified explicitly or not, this ONLY applies to the first project
given if multiple projects were given to this command; after the first project
is processed the value is set to 'none'.  The default account will only be
set if either the user does not currently have a default account (e.g. a new user)
or we are in clobber-default-account mode (see below).

The --[no]clobber-default-account flag determines whether the default account
computed above should be used to change the default account of the specified user
IF THAT USER already has a default account set (i.e. is not a new user of the cluster).  
If clobber-default-account is set, the user's default account will be set to the value
computed above (assuming a value was computed; i.e. did not request --default-account none),
otherwise the value computed above will only be used if the user does not already have
a default account.  The default is $tmp_clobber.

--highpriority is only used for condo accounts, and determines if user gets 
access to condo's hipriority QoS.  Default is no.

Dryrun mode just goes through the motions, verbose prints the sacctmgr/SQL
as executed.  Email will be sent to user and PoC of projects unless --noemail.

FIXME mode will attempt to repair an account; redoing all steps but not
complaining if already have been done.  Probably want --noemail when using it.

EOF
}

#------------------------------------------------------
#		Main
#------------------------------------------------------

my ($help, $dryrun, $verbose);
my $cluster;
my $projects=[];
my $users = [];
my ($reason);
my ($tmp, $where, $set, $error);
my $do_email=1;
my $highprio =0;
my $defacctflag;
my $FIXME;
my $clobber_defacct = $DEFAULT_CLOBBER_DEFACCT;

my $res = GetOptions(
	'help!' => \$help,
	'dryrun!' => \$dryrun,
	'verbose!' => \$verbose,
	'email!' => \$do_email,
	'fixme|FIXME!' => \$FIXME,
	'high-priority|high_priority|hi-priority|hi_priority|highpriority|hipriority!' 
		=> \$highprio,
	'clobber-default-account|clobber_default_account|cdefacct|clobber-defaultaccount|clobber_defaultaccount|clobber-defacct|clobber_defacct!' => \$clobber_defacct,

	'default-account|default_account|def-account|def_account|defaccount|defacct|defaultaccount:s' => \$defacctflag,

	'cluster=s' => \$cluster,
	'project=s' => $projects,
	'user=s' => $users,

	'reason=s' => \$reason,
);

unless ( $res )
{	usage_short;
	die "\nError processing arguments\n";
}

if ( $help )
{	usage_long;
	exit 0;
}

if (scalar(@ARGV) )
{	usage_short;
	$tmp = join ', ', @ARGV;
	die "\nUnrecognized arguments: [ $tmp ], aborting\n";
}

#Default cluster if needed
my $hpc = Glue::HPCC::Cluster->new_from_clustername($cluster);
die "\nInvalid cluster $cluster at " if $cluster && ! $hpc && ! ref($hpc);
die "\nError instantiating cluster $hpc " unless $hpc && ref($hpc);
$cluster = $hpc->db_clustername;
my $hpcUtil = $hpc->get_hpcUtil;
my $usrUtil = $hpcUtil->user;
#my $prjUtil = $hpcUtil->project;


$hpc->verbose($verbose);
$hpc->dryrun($dryrun);
Glue::HPCC::Utils::MailUtils::verbose($verbose);
Glue::HPCC::Utils::MailUtils::dryrun($dryrun);

unless ( scalar(@$projects) )
{	usage_short;
	die "\nNo projects given.\n";
}

unless ( scalar(@$users) )
{	usage_short;
	die "\nNo users given.\n";
}

#if ( defined $defacctflag && $defacctflag ne 'none' && $defacctflag ne 'no' )
#{	die "\nCannot use --default-account other than 'none' with multiple projects.\n"
#		if scalar(@$projects) > 1 ;
#}

if ( $defacctflag && $defacctflag eq 'no' )
{	$defacctflag = 'none';
}

my ($project, $proj);
my ($uname, $user );
my $defacct;
my @projs = ();
my ( $ptype, $ptypeid, $ptname );
my ( $projgrp, $projgid);

#Verify all specified projects exist and are active
foreach $project (@$projects)
{	my $prj = Glue::HPCC::Project::Project->project_by_name($hpc,
		$project);
	die "Invalid project '$project', not found, aborting.\n" unless $prj;
	die "Error looking up project $project: $prj at " unless ref($prj);

	die "Project '$project' is disabled; must enable it before " .
		"adding users.\n" unless $prj->active;

	push @projs, $prj;
	$projgrp = $prj->unixgroup;
	if ( $projgrp && $projgrp ne $usrUtil->default_users_unix_group_name )
	{	unless ( defined $projgid )
		{	$projgid = getgrnam($projgrp);
		}
	}

	$ptypeid = $prj->projtype;
	$ptype = Glue::HPCC::Project::ProjectType->projecttype_by_pkey(
		$hpc, $ptypeid);
	$ptname = $ptype->name;
	die "Cannot add members to project $project because is of type $ptname.\n"
		unless $ptype->can_have_members;
}


#Verify all specified users exist
foreach $uname (@$users)
{	if ( $uname ne lc $uname )
	{	warn "Replacing $uname with ", (lc $uname), "\n";
		$uname = lc $uname;
	}
	die "Invalid user $uname, not found in hesiod.\n"
		unless user_exists($uname);
}

#Do non-project specific work for the users
foreach $uname (@$users)
{	#IF creating home dir, etc., ALWAYS set GID to that of one of their projects
	#my $pwent = user_exists($uname);
	#my @pwent = split /:/, $pwent;
	#my $pgid = $pwent[3];

	my @fields = (
		hpc => $hpc,
		user => $uname,
		reason => $reason,
		FIXME => $FIXME,
	);

	#Set GID to that of first project with usable GID
	push @fields, gid => $projgid if defined $projgid;

	$error = $usrUtil->add_user_to_hpc_cluster( @fields);
	die "Error adding user $uname to cluster $cluster: $error.\n"
		if $error;
	unless ( $dryrun )
	{	#User won't necessarily be there after insert in dryrun mode
		$user = Glue::HPCC::Project::Person->person_by_username($hpc, 
			$uname);
		die "Error adding user $uname to cluster $cluster, " .
			"cannot find after insert at " 
			unless $user && ref($user);
		die "Error looking up user $uname: $user at " 
			unless ref($user);
	}
}

foreach $proj (@projs)
{	$project = $proj->name;

	#Set up stuff for emails.  Also validates project exists, etc.
	my $emails = Glue::HPCC::Project::PersonProject->get_poc_emails_for_project(
		$hpc, $project);
	die "Error getting emails for PoCs for $project: $emails at "
		unless $emails && ref($emails) eq 'ARRAY';

	$ptypeid = $proj->projtype;
	$ptype = Glue::HPCC::Project::ProjectType->projecttype_by_pkey($hpc, $ptypeid);
	$ptname = $ptype->name;

	#TODO: should use flags in project_type record
	my $accounts = [ $project ];
	if ( $ptname eq 'coop' )
	{	$accounts = [ $project, "${project}-hi" ];
	} elsif ( $ptype eq 'lease-hi' )
	{	$accounts = [ "${project}-hi" ];
	}
	if ( $highprio && $ptname ne 'condo' )
	{	warn "--high-priority flag ignored for project types other " .
			"than condo.\n";
		$highprio = 0;
	}

	$defacct = undef;
	if ( defined $defacctflag && $defacctflag ne 'none')
	{	$defacctflag = lc $defacctflag;
		if ( $defacctflag eq 'high' || $defacctflag eq 'hi' )
		{	die "Cannot set defacct to '$defacctflag' for project "
				. "of type $ptname.\n" 
				unless $ptype->has_high_priority;
			$defacct = $project . '-hi';
		} elsif ( $defacctflag eq 'std' || $defacctflag eq 'standard' )
		{	die "Cannot set defacct to '$defacctflag' for project "
				. "of type $ptname.\n" 
				unless $ptype->has_normal_priority;
			$defacct = $project;
		} elsif ( $defacctflag )
		{	die "Invalid value '$defacctflag' for --default-account ";
		} else
		{	$defacct = $project if $ptype->has_normal_priority;
			$defacct = $project . '-hi' 
				if $ptype->has_high_priority;
		}
		#IF multiple projects were given, do NOT set default account
		#on subsequent projects
		$defacctflag='none';
	} elsif ( ! defined $defacctflag )
	{	$defacct = $project if $ptype->has_normal_priority;
		$defacct = $project . '-hi' if $ptype->has_high_priority;
		#IF multiple projects were given, do NOT set default account
		#on subsequent projects
		$defacctflag='none';
	}


	my $tmplate = "${cluster}_welcome_new_user.tmplate";

	USER: foreach $uname (@$users)
	{	my @fields = (
			user => $uname,
			project => $project,
			reason => $reason,
			FIXME => $FIXME,
		);

		my $defaccount;
		$defaccount = $usrUtil->get_default_account_for_user(
				hpc => $hpc, user => $uname,
				nocroak => 1, nowarn => 1);
		#print STDERR "Started with defaccount=$defaccount\n";

		if ( $defacct )
		{	#A default account was requested
			#But only add if user has none OR --clobber-default-account
			if ( $clobber_defacct || ! $defaccount )
			{	push @fields, defaultaccount => $defacct;
			}
		}
		push @fields, highpriority => $highprio if $ptname eq 'condo';

		$error = $usrUtil->add_user_to_project_in_cluster( 
			hpc=>$hpc, @fields);
		if ( $error )
		{	warn "Error adding user $uname to $project in " .
				"$cluster: $error\n";
			warn "No mail sent for user $uname.\n";
			next USER;
		}

		#Verify user exists and get default account
		if ( $dryrun )
		{	#Fake out to avoid errors in dryrun mode
			#This might not be accurate if multiple projects given
			$defaccount = $project;
			$defaccount = $defacct if $defacct;
		} else
		{	$defaccount = $usrUtil->get_default_account_for_user(
				hpc => $hpc, user => $uname);
		}
		my $user = Glue::HPCC::Project::Person->person_by_username(
			$hpc, $uname);
		die "Error: cannot find user $uname in PDB after adding, at "
			unless $user || $dryrun;
		if ( $user && ! ref($user) )
		{ 	warn "Error looking up user $uname in PDB: $user at ";
			warn "No mail sent for user $uname.\n";
			next USER;
		}
		my ($email, $full);
		if ( $dryrun && ! $user )
		{	#Special case, fake lookup of user
			$email = $uname . '@umd.edu';
			$full = "fullname of $uname";
		} else
		{	$email = $user->email;
			$full = $user->fullname;
		}
		my $cclist = [ @cc_on_everything, @$emails ];
		#Don't cc PoC when adding PoC
		my %temp = map { $_ => $_ } @$cclist;
		delete $temp{$email}; 
		$cclist = [ values %temp ];
	
		#Commit after each user added
		$hpc->db_commit;
		#Send email.  Currently a bit ugly, an email per user per project
		#Would be nicer if could combine, but since usually only adding one or two
		#users at a time to a single project, not putting much effort into it
		my $tthash = 
		{	USERNAME => $uname,
			EMAIL => $email,
			FULLNAME => $full,
			PROJECT => $project,
			DEFACCOUNT => $defaccount,
			ACCOUNTS => $accounts,
			CCLIST => $cclist,
			CLUSTER => $cluster,
		};
		send_template_email($tmplate, $tthash) if $do_email;
			
	}
}
