#!/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.5'
author  = 'Walter Brisken and David Gordon'
verdate = '20140105'
# 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
		
def testdifxversion(kv, override):
	difxVersion = getenv('DIFX_VERSION')
	difxLabel = getenv('DIFX_LABEL')
	fitsstop = False
	if difxVersion == None:
		print 'Error: env var DIFX_VERSION is not set!'
		exit(0)
	elif difxVersion != kv['DiFX']:
		if override:
			print 'Overriding version mismatch: %s != %s' % (difxVersion, kv['DiFX'])
			fitsstop = False
			kv['DiFX'] = difxVersion
		else:
			print 'Error: FITS file DiFX version = %s' % kv['DiFX']
			print 'and environment DiFX version = %s' % difxVersion
			fitsstop = True

	if difxLabel != None:
		if not 'label' in kv:
			if override:
				print 'Override causing DiFX label to be set to %s' % difxLabel
				kv['label'] = difxLabel
				difxVersion = 'unknown'
			else:
				print 'Error: file DiFX label is not set but environment DiFX label is %s' % difxLabel
				exit(0)
		if difxLabel != kv['label']:
			if override:
				print 'Overriding DiFX label mismatch: %s != %s' % (difxLabel, kv['label'])
				kv['label'] = difxLabel
	elif 'label' in kv:
		if override:
			print 'Override causing DiFX label to be unset (it was %s)' % kv['label']
			del kv['label']
		else:
			print 'Error: file DiFX label is set to %s but environment DiFX label not set' % kv['label']
			exit(0)

	return fitsstop

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):
		return testdifxversion(self.kv, override)

	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):
		testdifxversion(self.kv, override)

	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 fitslist 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 ''