#!/usr/bin/env python

from sys import exit, argv, stdout
from os import getenv, system, stat, popen, getcwd
from string import split, strip, upper
from getopt import getopt
from os.path import isfile, isdir
from time import gmtime

program = 'difxarch'
version = '0.4'
author  = 'Walter Brisken and David Gordon'
verdate = '20120718'
# Modified by David Gordon to archive Mark4 files from makemark4.

def usage(prog):
	print '%s ver %s  %s  %s' % (program, version, author, verdate)
	print '\nUsage:  %s [options] <passName> [<passName2> [<passName3> ... ] ]' % prog
	print '\noptions can include:\n'
	print '  --verbose'
	print '  -v         Send more output to the screen\n'
	print '  --help'
	print '  -h         Print this help information and quit\n'
	print '  --override-version'
	print '             To force operation with mixed DiFX versions\n'
	print '  --pretend'
	print '  -p         Pretend mode: don\'t actually do anything\n'
	print '  --fits '
	print '  -f         Process only fits files\n'
	print '  --mark4'
	print '  -m         Process only mark4 files\n'
	exit(0)

def teste2ecopy(pretend):
	d = popen('which e2ecopy', 'r').readlines()
	if len(d) < 1:
		print '\nError: program e2ecopy not found.  The difx_db package'
		print 'Is not properly installed and difxarch cannot continue.'
		if not pretend:
			exit(0)
	e = popen('ls -l %s' % d[0], 'r').readlines()
	s = split(e[0])
	if s[2] != 'root':
		print '\nError: program e2ecopy is not owned by root.  Contact'
		print 'a sysadmin to have ownership changed to root and have'
		print 'the +s bit set on its permissions.'
		if not pretend:
			exit(0)
	if s[0][3] != 's':
		print '\nError: program e2ecopy does not have +s permissions.'
		print 'Have a sysadmin run chmod +s on the appropriate copy of e2ecopy'
		if not pretend:
			exit(0)

def mjd2gmt(mjd):
	return 86400.0*(mjd - 40587.0)

def genDateStr(t):
	d = gmtime(t)
	return '%02d%02d%02dT%02d%02d%02d' % \
		(d.tm_year % 100, d.tm_mon, d.tm_mday, \
		d.tm_hour, d.tm_min, d.tm_sec)

def parsekv(str):
	kv = {}
	ss = split(str)
	for s in ss:
		p = split(s, '=')
		if len(p) != 2:
			print 'Error parsing key=value statement: %s', s
			return {}
		kv[p[0]] = p[1]
	return kv

class FitsList:
	def __init__(self, filename):
		self.filename = filename
		self.files = []
		data = open(filename).readlines()
#		if len(data) < 2:
#			print 'Malformed .fitslist file %s' % filename
#			exit(0) 
		self.kv = parsekv(data[0])
#		print '\n kv = %s ' % self.kv
#		if len(self.kv) < 1:
#			print 'Malformed .fitslist file %s line 1' % filename
#			exit(0)
		for d in data[1:]:
			s = split(d)
			if s < 2:
				continue;
			self.files.append(s[0])
		
	def testdata(self):
		if len(self.files) < 1:
			print 'Malformed .fitslist file %s' % self.filename
			fitsstop = True
			return fitsstop
		else:
			fitsstop = False

		if len(self.kv) < 1:
			print 'Malformed .fitslist file %s line 1' % self.filename
			fitsstop = True
			return fitsstop
		else:
			fitsstop = False
		return fitsstop

	def testversion(self, override):
		difxVersion = getenv('DIFX_VERSION')
		fitsstop = False
		if difxVersion == None:
			print 'Warning: env var DIFX_VERSION is not set!'
			if override:
				difxVersion = 'unknown'
				fitsstop = False
			else:
				fitsstop = True
				exit(0)
		elif difxVersion != self.kv['DiFX']:
			if override:
				print 'Overriding version mismatch: %s != %s' % \
					(difxVersion, self.kv['DiFX'])
				difxVersion = 'unknwown'
				fitsstop = False
			else:
				print 'Error: FITS file DiFX version = %s' % self.kv['DiFX']
				print 'and current DiFX version = %s' % difxVersion
				fitsstop = True

		self.kv['DiFX'] = difxVersion
                return fitsstop

	def verify(self, queuedir):
 		dir = queuedir + '/' + self.kv['exper']
		if not isdir(dir):
			print 'Purported project queue directory %s not found' % dir
			return False
		r = True
		for f in self.files:
			fn = dir + '/' + f
			if not isfile(fn):
			        r = False
			        if verbose: 
		     	        	print 'FITS file %s not found' % fn
		return r
				

	def show(self, indent=0):
		id = ' '*indent
		print '%sFitsList: %s' % (id, self.filename)
		id = ' '*(indent+2)
		for key in self.kv.keys():
			print '%sKV: %s = %s' % (id, key, self.kv[key])
		for j in self.files:
			print '%sFILE: %s' % (id, j)

	def archive(self, queuedir, archdir, verbose, pretend):
		cwd = getcwd()

 		indir = queuedir + '/' + self.kv['exper']
		outdir = archdir + '/' + self.kv['exper'] + '.' + self.kv['pass']

                if verbose:
                        print '\n Fits outdir: %s \n' % outdir

		if isdir(outdir):
			print 'Error: FITS archive staging directory %s already exists.' % outdir
			print 'Please try again later after previous arching completes.'
			if not pretend:
				return

		if verbose:
			options = '-v'
		else:
			options = ''

		cmd = 'e2ecopy %s %s %s' % (options, indir, outdir)
		for fn in self.files:
			cmd += ' %s' % fn
			if fn[-8:] == '.idifits':
				fn2 = fn[0:-8] + '.jobmatrix.txt'
				if isfile(fn2):
					cmd += ' %s' % fn2

		difxlogFile = self.kv['pass'] + '.difxlog.gz'
		if isfile(difxlogFile):
			datestr = genDateStr(mjd2gmt(float(self.kv['mjd'])))
			e2eFile = '%s/%s:VLBA_%s_%s_%s.difx.log.gz' % (cwd, difxlogFile, upper(self.kv['exper']), self.kv['pass'], datestr)
			cmd += ' ' + e2eFile
		elif verbose:
			print 'Cannot find %s so won\'t archive it.' % difxlogFile

                print '\n Now archiving FITS files.\n '

		if verbose:
			print 'Executing: %s' % cmd

		if pretend:
			print '\n**************************************************'
			print 'Pretend mode activated; not actually doing anything!'
			print '**************************************************\n'
			return

		system(cmd)

		nErr = 0
		for fn in self.files:
			if not isfile('%s/%s' % (outdir, fn)):
				print 'Error: File %s did not copy!' % fn
				nErr += 1
			elif stat('%s/%s' % (indir, fn)).st_size != \
			     stat('%s/%s' % (outdir, fn)).st_size:
				print 'Error: Incomplete copy of %s!' % fn

		if nErr > 0:
			print '\n*************************************\n'
			print '%d of %d FITS files had copy errors' % (nErr, len(fn))
			print 'Original files are not being removed!'
			print '\n*************************************\n'
		else:
			cmd = 'rm -f'
			for fn in self.files:
				cmd += ' %s/%s' % (indir, fn)
			if verbose:
				print 'Executing: %s' % cmd
			system(cmd)


class Mark4List:
	def __init__(self, filename):
		self.filename = filename
		self.files = []
		data = open(filename).readlines()
		if len(data) < 2:
			print 'Malformed .mark4list file %s' % filename
			exit(0)
		self.kv = parsekv(data[0])
		if len(self.kv) < 1:
			print 'Malformed .mark4list file %s line 1' % filename
			exit(0)
		for d in data[1:]:
			s = split(d)
			if s < 2:
				continue;
			self.files.append(s[0])
		
	def testversion(self, override):
		difxVersion = getenv('DIFX_VERSION')
		if difxVersion == None:
			print 'Warning: env var DIFX_VERSION is not set!'
			if override:
				difxVersion = 'unknown'
			else:
				exit(0)
		elif difxVersion != self.kv['DiFX']:
			if override:
				print 'Overriding version mismatch: %s != %s' % \
					(difxVersion, self.kv['DiFX'])
				difxVersion = 'unknwown'
			else:
				print 'Error: mark4 file DiFX version = %s' % self.kv['DiFX']
				print 'and current DiFX version = %s' % difxVersion
				exit(0)

		self.kv['DiFX'] = difxVersion

	def verify(self, queuedir):
 		dir = queuedir + '/' + self.kv['exper']

		if not isdir(dir):
			print 'Purported project queue directory %s not found' % dir
			return False
		r = True
		for f in self.files:
			fn = dir + '/' + f
			if not isfile(fn):
				print 'Mark4 file %s not found' % fn
				r = False
		return r
				

	def show(self, indent=0):
		id = ' '*indent
		print '%sMark4List: %s' % (id, self.filename)
		id = ' '*(indent+2)
		for key in self.kv.keys():
			print '%sKV: %s = %s' % (id, key, self.kv[key])
		for j in self.files:
			print '%sFILE: %s' % (id, j)

	def archive(self, queuedir, archdir, verbose, pretend):
		cwd = getcwd()

 		indir = queuedir + '/' + self.kv['exper']

		outdir = archdir + '/' + self.kv['exper'] + '.' + self.kv['pass'] \
		        + 'mark4'

                if verbose: 
                        print '\n outdir: %s \n' % outdir

		if isdir(outdir):
			print 'Error: archive staging directory %s already exists.' % outdir
			print 'Please try again later after previous arching completes.'
			if not pretend:
				exit(0)

		if verbose:
			options = '-v'
		else:
			options = ''

		cmd = 'e2ecopy %s %s %s' % (options, indir, outdir)
		for fn in self.files:
			cmd += ' %s' % fn

# Add mark4 metadata file
		metadatafile = self.kv['exper'] + '.mark4.metadata.txt'
		if isfile(indir + '/' + metadatafile):
		        cmd += ' %s' % metadatafile
		else: 
		        print '\n Cannot find metadata file %s !!!!! \n' % metadatafile


                print '\n Now archiving mark4 files.\n '

		if verbose:
			print 'Executing: %s' % cmd

		if pretend:
			print '\n**************************************************'
			print 'Pretend mode activated; not actually doing anything!'
			print '**************************************************\n'
			return

		system(cmd)

		nErr = 0
		for fn in self.files:
			if not isfile('%s/%s' % (outdir, fn)):
				print 'Error: File %s did not copy!' % fn
				nErr += 1
			elif stat('%s/%s' % (indir, fn)).st_size != \
			     stat('%s/%s' % (outdir, fn)).st_size:
				print 'Error: Incomplete copy of %s!' % fn

		if nErr > 0:
			print '\n*************************************\n'
			print '%d of %d mark4 files had copy errors' % (nErr, len(fn))
			print 'Original files are not being removed!'
			print '\n*************************************\n'
		else:
			cmd = 'rm -f'
			for fn in self.files:
				cmd += ' %s/%s' % (indir, fn)
			if verbose:
				print 'Executing: %s' % cmd
			system(cmd)


# main below here

print ''

archroot = getenv('DIFX_ARCHIVE_ROOT')
if archroot == None:
	print 'Error: Environment variable DIFX_ARCHIVE_ROOT not set.'
	print 'Cannot proceed.\n'
	exit(0)

queueroot = getenv('DIFX_QUEUE_BASE')


if queueroot == None:
	print 'Error: Environment variable DIFX_QUEUE_BASE not set.'
	print 'Cannot proceed.\n'
	exit(0)

optlist, args = getopt(argv[1:], 'hvpfm', ['help', 'verbose', 'pretend', 'fits', 'mark4', 'override-version'])

overrideVersion = False
verbose = False
pretend = False
stop = False
domark4 = True
dofits = True
doofits = True

for o, a in optlist:
	if o == '--override-version':
		overrideVersion = True
	elif o in ('-v', '--verbose'):
		verbose = True
	elif o in ('-p', '--pretend'):
		pretend = True
	elif o in ('-h', '--help'):
		usage(argv[0])
	elif o in ('-f', '--fits'):
		domark4 = False
	elif o in ('-m', '--mark4'):
		doofits = False
	else:
		print 'Error: unknown command line option: %s' % o
		stop = True

if not doofits and not domark4:
        print '\n Both doofits and domark4 are False, therfore I am quitting!\n'
        stop = True

if len(args) < 1:
	print 'Error: no input files provided'
	stop = True

if stop:
	print '\nRun with -h for help information.\n'
	exit(0)

teste2ecopy(pretend)

if doofits:
   for arg in args:
	if arg[-9:] == '.fitslist':
		fitslistfile = arg
	elif arg[-10:] == '.mark4list':
	        sfits = split(arg,'.')
		fitslistfile = sfits[0] + '.fitslist'
	else:
		fitslistfile = arg + '.fitslist'

	if not isfile(fitslistfile):
		print '\n File %s not found, no FITS archiving. \n' % fitslistfile
		dofits = False
		if not domark4:
                       exit(0)
        if dofits:
	        fl = FitsList(fitslistfile)

        if dofits:
	        fitsstop = fl.testdata()
		if fitsstop:
		       dofits = False

        if dofits:
	        fitsstop = fl.testversion(overrideVersion)
		if fitsstop:
		       dofits = False

        if dofits:
	        if verbose:
		        fl.show()
	
        if dofits:
	        OK = fl.verify(queueroot)

        if dofits:
	        if not OK:
		        print '\n fits list verification failed! Will not archive FITS files.'
			dofits = False
			if not domark4:
                           exit(0)
	
        if dofits:
	        OK = fl.archive(queueroot, archroot, verbose, pretend)
	else:
	        print '\n Not archiving any FITS files. \n'



if domark4:
   for arg in args:
	if arg[-9:] == '.fitslist':
	        smark4 = split(arg,'.')
		mark4listfile = smark4[0] + '.mark4list'
	elif arg[-10:] == '.mark4list':
		mark4listfile = arg
	else:
		mark4listfile = arg + '.mark4list'

	if not isfile(mark4listfile):
		print '\n File %s not found, no mark4 archiving. \n' % mark4listfile
		exit(0)

	ml = Mark4List(mark4listfile)

	ml.testversion(overrideVersion)

	if verbose:
		ml.show()
	
	OK = ml.verify(queueroot)

	if not OK:
		print '\n mark4 list verification failed! Will not archive mark4 files. \n'
		exit(0)
	
	OK = ml.archive(queueroot, archroot, verbose, pretend)

print ''