#!/usr/local/bin/python

import sys, os
import re

import tempfile
import subprocess
import logging 
import time
import signal

import util

from ssh_emulab_segment import ssh_emulab_segment
from xmlrpc_emulab_segment import xmlrpc_emulab_segment


class start_segment(ssh_emulab_segment, xmlrpc_emulab_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None,
	    ops=None, cert=None):
	ssh_emulab_segment.__init__(self, log=log, keyfile=keyfile, debug=debug)
	xmlrpc_emulab_segment.__init__(self, boss=boss, ops=ops, cert=cert)

    def set_up_experiment_filespace(self, user, host, pid, eid, tmpdir):
	"""
	Send all the software and configuration files into the experiment's
	file space.  To reduce the number of ssh connections, we script many
	changes and execute the script.
	"""
	# Configuration directories on the remote machine
	proj_dir = "/proj/%s/exp/%s/tmp" % (pid, eid)
	softdir_root = "/proj/%s/software/" % pid
	softdir = "%s/%s" % (softdir_root, eid)
	# Local software dir
	lsoftdir = "%s/software" % tmpdir

	# Open up a temporary file to contain a script for setting up the
	# filespace for the new experiment.
	self.log.info("[start_segment]: creating script file")
	try:
	    sf, scriptname = tempfile.mkstemp()
	    scriptfile = os.fdopen(sf, 'w')
	except EnvironmentError:
	    return False

	scriptbase = os.path.basename(scriptname)

	# Script the filesystem changes
        
	# Set umask to 007.  Since we are operating
        # within /proj, we want to permit all proj users
        # read and write.  This is especially necessary
	# since creating directory hierarchy (mkdir -p).
	print >>scriptfile, "umask 007"
	print >>scriptfile, "/bin/rm -rf %s" % proj_dir
	# Clear and create the software directory
	print >>scriptfile, "/bin/rm -rf %s/*" % softdir
	print >>scriptfile, 'mkdir %s' % proj_dir
	print >>scriptfile, 'chmod 0770 %s' % proj_dir
	if os.path.isdir(lsoftdir):
	    print >>scriptfile, 'mkdir %s' % softdir_root
	    print >>scriptfile, 'chmod 0770 %s' % softdir_root
	    print >>scriptfile, 'mkdir %s' % softdir
	    print >>scriptfile, 'chmod 0770 %s' % softdir
	print >>scriptfile, "rm -f %s" % scriptbase
	scriptfile.close()

	# Move the script to the remote machine
	# XXX: could collide tempfile names on the remote host
	if self.scp_file(scriptname, user, host, scriptbase):
	    os.remove(scriptname)
	else:
	    return False

	# Execute the script (and the script's last line deletes it)
	if not self.ssh_cmd(user, host, "sh -x %s" % scriptbase):
	    return False

	for f in os.listdir(tmpdir):
	    if not os.path.isdir("%s/%s" % (tmpdir, f)):
		if not self.scp_file("%s/%s" % (tmpdir, f), user, host,
			"%s/%s" % (proj_dir, f)):
		    return False
	if os.path.isdir(lsoftdir):
	    for f in os.listdir(lsoftdir):
		if not os.path.isdir("%s/%s" % (lsoftdir, f)):
		    if not self.scp_file("%s/%s" % (lsoftdir, f), 
			    user, host, "%s/%s" % (softdir, f)):
			return False
	return True


    def __call__(self, parent, eid, pid, user, tclfile, tmpdir, timeout=0, 
	    gid=None):
	"""
	Start a sub-experiment on a federant.

	Get the current state, and terminate the experiment if it exists. The
	group membership of the experiment is difficult to determine or change,
	so start with a clean slate.  Create a new one and ship data
	and configs and start the experiment.  There are small ordering
	differences based on the initial state of the sub-experiment.
	"""

	state = self.get_state(pid, eid)

	if state != 'none':
	    self.terminate_exp(pid, eid)

	if not self.make_null_experiment(pid, eid, tmpdir, gid):
	    return False

	if not self.set_up_experiment_filespace(user, self.ops, 
		pid, eid, tmpdir):
	    return False

	# Put the file into a string to pass to emulab.
	try:
	    tcl = "".join([ l for l in open(tclfile,"r")])
	except EnvironmentError, e:
	    self.log.error("Can't read %s: %s" % (tclfile, e))
	    return False
	
	# Stage the new configuration
	if not self.modify_exp(pid, eid, tcl):
	    self.log.error("modify failed")
	    return False

	if not self.swap_exp(pid, eid, 'in'):
	    self.log.error("swap in failed")
	    return False
	# Everything has gone OK.
	self.get_mapping(pid,eid)
	return True

class stop_segment(ssh_emulab_segment, xmlrpc_emulab_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None, 
	    ops=None, cert=None):
	ssh_emulab_segment.__init__(self, log=log, keyfile=keyfile, debug=debug)
	xmlrpc_emulab_segment.__init__(self, boss=boss, ops=ops, cert=cert)

    def __call__(self, parent, user, pid, eid, gid=None, terminate=False):
	"""
	Stop a sub experiment by calling swapexp on the federant
	"""

	self.log.info("[stop_segment]: Stopping %s" % eid)
	rv = False
	try:
	    # Clean out tar files: we've gone over quota in the past
	    self.ssh_cmd(user, self.ops, 
		    "rm -rf /proj/%s/software/%s" % (pid, eid))
	    rv = self.swap_exp(pid, eid, 'out')
	    if terminate:
		rv = self.terminate_exp(pid, eid)
	except self.ssh_cmd_timeout:
	    rv = False
	return rv

class info_segment(ssh_emulab_segment, xmlrpc_emulab_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None, 
	    ops=None, cert=None):
	ssh_emulab_segment.__init__(self, log=log, keyfile=keyfile, debug=debug)
	xmlrpc_emulab_segment.__init__(self, boss=boss, ops=ops, cert=cert)

    def __call__(self, parent, user, pid, eid):
	self.log.info("[info_segment]: Getting info from %s" % eid)
	self.get_mapping(pid,eid)
	return True

class operation_segment(ssh_emulab_segment, xmlrpc_emulab_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None, 
	    ops=None, cert=None):
	ssh_emulab_segment.__init__(self, log=log, keyfile=keyfile, debug=debug)
	xmlrpc_emulab_segment.__init__(self, boss=boss, ops=ops, cert=cert)

    def __call__(self, parent, op, targets, param, top):
	for l, p in targets.items():
	    self.log.info("[operation_segment]: Calling op %s on %s(%s)" % \
		    (op, l,p))
	    self.do_operation(op, l, p, param, top)
	return True