#!/usr/bin/env python3

#**************************************************************************
#   Copyright (C) 2008-2019 by Walter Brisken                             *
#                                                                         *
#   This program is free software; you can redistribute it and/or modify  *
#   it under the terms of the GNU General Public License as published by  *
#   the Free Software Foundation; either version 3 of the License, or     *
#   (at your option) any later version.                                   *
#                                                                         *
#   This program is distributed in the hope that it will be useful,       *
#   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
#   GNU General Public License for more details.                          *
#                                                                         *
#   You should have received a copy of the GNU General Public License     *
#   along with this program; if not, write to the                         *
#   Free Software Foundation, Inc.,                                       *
#   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
#**************************************************************************
#===========================================================================
# SVN properties (DO NOT CHANGE)
#
# $Id$
# $HeadURL: $
# $LastChangedRevision$
# $Author$
# $LastChangedDate$
#
#============================================================================

# Note: this utility can run under python2.7 or python3

from sys import argv, exit
from os import popen, getenv
from threading import Thread
from glob import glob

program = 'listmodules'
author  = 'Walter Brisken'
version = '0.1'
verdate = '20191007'

mk5list = 'mk5list'

def usage():
	print('\n%s ver. %s  %s %s' % (program, version, author, verdate))
	print('\nA program to find required Mark5 modules and write the machines file')
	print('appropriate for a particular DiFX job.')
	print('\nUsage : %s [options]' % argv[0])
	print('\noptions can include:')
	print('\n  -h or --help')
	print('          print this usage info and exit')
	print('\n  -m <machinesfile> or --machines <machinesfile>')
	print('          use <machinesfile> instead of $DIFX_MACHINES')
	print('\nThis program responds to the following environment variables:\n')
	print('DIFX_MACHINES must point to the machines file if no <machinesfile> is specifided.\n')
	exit(1)


def readmachines(machinesfile):
	machines = []
	cores = []
	ismk5 = {}
	lines = open(machinesfile).readlines()
	for l in lines:
		l = l.strip().split('#')[0]
		s = l.split()
		if len(s) >= 2:
			machines.append(s[0])
			cores.append(int(s[1]))
			if s[0][:5] == 'mark5':
				ismk5[s[0]] = 1
			else:
				ismk5[s[0]] = 0
		if len(s) >= 3:
			ismk5[s[0]] = int(s[2])
	return machines,cores,ismk5

def gethostname():
	o = popen('hostname', 'r')
	host = o.readline().strip()
	o.close()
	return host

class querythread(Thread):
	def __init__(self, machine):
		Thread.__init__(self)
		self.machine = machine
		self.result = ''
	def run(self):
		cmd = 'ssh %s %s 2>&1' % (self.machine, mk5list)
		o = popen(cmd, 'r') 
		lines = o.readlines()
		o.close()
		self.result = lines[-1].strip()

def mk5query(machines, ismk5):
	qlist = []
	results = []
	for m in machines:
		if ismk5[m]:
			qt = querythread(m)
			qlist.append(qt)
			qt.start()
	for qt in qlist:
		qt.join()
		results.append([qt.machine, qt.result])
	return results


def run(machinesfile):
	machines,cores,ismk5 = readmachines(machinesfile)

	results = mk5query(machines, ismk5)
	results.sort()

	for r in results:
		print(r[0], r[1])

machinesfile = getenv('DIFX_MACHINES')

a = 1
while a < len(argv):
	arg = argv[a]
	if arg == '-h' or arg == '--help':
		usage()
	elif arg == '-m' or arg == '--machinesfile':
		a += 1
		if a >= len(argv):
			print('No machinesfile specified!')
			exit(1)
		machinesfile = argv[a]
	a += 1

quit = False
	
if machinesfile == None:
	print('DIFX_MACHINES env var not defined!')
	quit = True

elif len(glob(machinesfile)) != 1:
	print('Machinesfile %s not found.' % machinesfile)
	quit = True
	
if quit:
	print('%s quitting.' % program)
	exit(1)
	
v = run(machinesfile)