import os
import sys
import time
import socket
import threading
import subprocess
import sunhpc.commands
class Parallel(threading.Thread):
    def __init__(self, cmdclass, cmd, host, hostif, stats, collate):
        threading.Thread.__init__(self)
        self.cmd = cmd
        self.host = host
        self.hostif = hostif
        self.stats = stats
        self.collate = collate
        self.cmdclass = cmdclass

    def run(self):
        starttime = time.time()
        self.p = subprocess.Popen(self.cmd,
            stdin = subprocess.PIPE, stdout = subprocess.PIPE,
            stderr = subprocess.STDOUT)

        for line in self.p.stdout.readlines():
            if self.collate:
                self.cmdclass.addOutput(self.host, line[:-1])
            else:
                print (line[:-1].decode('utf-8'))

        if self.stats:
            msg = 'command on host %s took %f seconds' % \
                (self.host, time.time() - starttime)

            if self.collate:
                self.cmdclass.addOutput(self.host, msg)
            else:
                print (msg.decode('utf-8'))

    def kill(self):
        os.kill(self.p.pid, 9)
    
class command(sunhpc.commands.HostArgumentProcessor,
        sunhpc.commands.run.command):
    MustBeRoot = 0
    
class Command(command):
    """
    Run a command for each specified host.

    <arg optional='1' type='string' name='host' repeat='1'>
    Zero, one or more host names. If no host names are supplied, the command
    is run on all 'managed' hosts. By default, all compute nodes are
    'managed' nodes. To determine if a host is managed, execute:
    'sunhpc list host attr hostname | grep managed'. If you see output like:
    'compute-0-0: managed true', then the host is managed.
    </arg>

    <arg type='string' name='command'>
    The command to run on the list of hosts.
    </arg>

    <param type='boolean' name='managed'>
    Run the command only on 'managed' hosts, that is, hosts that generally
    have an ssh login. Default is 'yes'.
    </param>

    <param type='boolean' name='x11'>
    If 'no', disable X11 forwarding when connecting to hosts.
    Default is 'yes'.
    </param>

    <param type='string' name='timeout'>
    Sets the maximum length of time (in seconds) that the command is
    allowed to run.
    Default is '30'.
    </param>

    <param type='string' name='delay'>
    Sets the time (in seconds) to delay between each executed command
    on multiple hosts. For example, if the command is run on two
    hosts and if the delay is 10, then the command will be executed on host
    1, then 10 seconds later, the command will be executed on host 2.
    Default is '0' (no delay).
    </param>

    <param type='string' name='stats'>
    Display performance statistics if this parameter is set to 'yes'.
    Default is 'no'.
    </param>

    <param type='string' name='collate'>
    Prepend the hostname to every output line if this parameter is set to
    'yes'.
    Default is 'no'.
    </param>

    <param type='string' name='num-threads'>
    The number of threads to start in parallel. If num-threads is 0, then
    try to run the command in parallel on all hosts. Default is '128'.
    </param>

    <param type='string' name='command'>
    Can be used in place of the 'command' argument.
    </param>

    <example cmd='run host compute-0-0 command="hostname"'>
    Run the command 'hostname' on compute-0-0.
    </example>

    <example cmd='run host compute "ls /tmp"'>
    Run the command 'ls /tmp/' on all compute nodes.
    </example>
    """

    def nodeup(self, hostif):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(2.0)
        try:
            sock.connect((hostif, 22))
            buf = sock.recv(64)
            sock.send('SSH-2.0-nodeup\r\n')
            buf = sock.recv(1024)
            sock.close()
        except:
            return 0
        return 1

    def run(self, params, args):
        (args, command) = self.fillPositionalArgs(('command', ))

        if not command:
            self.abort('must supply a command')

        (managed, x11, t, d, s, c, n) = \
            self.fillParams([
                ('managed', 'y'),
                ('x11', 'y'),
                ('timeout', '30'),
                ('delay', '0'),
                ('stats', 'n'),
                ('collate', 'n'),
                ('num-threads', '8')
            ])

        try:
            timeout = int(t)
        except:
            self.abort('"timeout" must be an integer')

        if timeout < 0:
            self.abort('"timeout" must be a postive integer')

        try:
            numthreads = int(n)
        except:
            self.abort('"num-threads" must be an integer')

        try:
            delay = float(d)
        except:
            self.abort('"delay" must be a floating point number')

        hosts = self.getHostnames(args, self.str2bool(managed))
        # This is the same as doing -x using ssh.  Might be useful
        # for the common case, but required for the Viz Roll.
        if not self.str2bool(x11):
            try:
                del os.environ['DISPLAY']
            except KeyError:
                pass

        collate = self.str2bool(c)
        stats = self.str2bool(s)
        if collate:
            self.beginOutput()

        if numthreads <= 0:
            numthreads = len(hosts)

        threads = []
        i       = 0
        work    = len(hosts)
        while work:
            localhost = socket.gethostname().split('.')[0]
            while i < numthreads and i < len(hosts):
                host = hosts[i]
                # Is this host me?
                runlocal = (localhost == host.split('.')[0])
                i += 1  

                try:
                    hnet=self.db.getHostAttr(host,'primary_net')
                    query="""select net.ip from networks net, nodes n, subnets s where 
                            net.node=n.id and net.subnet=s.id and 
                            n.name='%s' and s.name='%s' """ % (host, hnet)
                    self.db.execute(query)
                    hostif,=self.db.fetchone()
                except:
                    hostif=host
                #
                # first test if the node is up and responding
                # to ssh
                #
                if not runlocal and not self.nodeup(hostif):
                    if collate:
                        self.addOutput(host, 'down')
                    else:
                        print ('%s: down' % host)

                    numthreads += 1
                    work -= 1
                    continue
                #
                # fire off the command
                #
                if runlocal:
                    cmd = ('bash', '-c', command)
                else:
                    cmd = ('ssh', hostif, command)

                p = Parallel(self, cmd, host, hostif, stats, collate)
                p.start()
                threads.append(p)

                if delay > 0:
                    time.sleep(delay)
            #
            # collect completed threads
            #
            try:
                totaltime = time.time()
                while timeout == 0 or \
                    (time.time() - totaltime) < timeout:

                    active = threading.enumerate()

                    t = threads
                    for thread in t:
                        if thread not in active:
                            thread.join(0.1)
                            threads.remove(thread)
                            numthreads += 1
                            work -= 1

                    if len(active) == 1:
                        break
                    #
                    # don't burn a CPU while waiting for the
                    # threads to complete
                    #
                    time.sleep(0.5)

            except KeyboardInterrupt:
                #
                # try to collect all the active threads
                #
                active = threading.enumerate()

                t = threads
                for thread in t:
                    if thread not in active:
                        thread.join(0.1)
                        threads.remove(thread)
                #
                # no more work to do if the user hits
                # control-c
                #
                work = 0
        #
        # kill all still active threads
        #
        active = threading.enumerate()
        if len(active) >= 2:
            for i in range(1, len(active)):
                active[i].kill()

        if collate:
            self.endOutput(padChar='')

RollName = "base"
