
"""Give a list of dirac prod id and it creates folder and download the log associated."""


"""
Modified by Jackson Smith,  Feb 2014, to account for changes in the data structure since
the script was first written.

Hacked together by Alex Shires, using Sam Gregson's script XMLMerge.py ( from Rob Lambert) 
and using Marc-Olivier's original gauss statistics script as the base

requires a modified version of dirac-production-download-logs-without-check.py

requires XMLMerge.py
"""

import sys, string
from subprocess import Popen, PIPE, STDOUT
from tempfile import NamedTemporaryFile
from glob import glob
import os
from datetime import date
from optparse import OptionParser
import shutil
import tarfile




#dirac_download_script = 'dirac-production-download-files-withoutcheck.py'
dirac_download_script = 'dirac-production-download-logs-withoutcheck.py'

castor_location = '/castor/cern.ch/grid/lhcb/backup/log/'

def parse_options():
    
    opt = OptionParser("usage: %prog [options] <ProdIDs>\n\n       <ProdIDs> : list of ProdID(s), comma-separated, with no blank spaces.")


    opt.add_option("-n", "--number-of-logs",
                   action="store", type="int", dest="nb_logs",
                   default=300,
                   metavar='<NB_LOGS>',
                   help="number of logs to download [default : %default]")

    opt.add_option("-v", "--version",
                   action="store", type="string", dest="MCversion",
                   default="2012",
                   metavar='<VERSION>',
                   help="version of MC [default : %default]")

    opt.add_option("--delta-size",
                   action="store", type="float", dest="delta",
                   default=0.07,
                   metavar='<DELTA>',
                   help="log files smaller by <DELTA> from largest file of sample will be deleted [default : %default]")
    
    opt.add_option("--usage",
                   action="store_true", dest="print_usage",
                   default=False,
                   help="show this help message and exit")
 
                  
    opt.add_option("-c","--castor",action="store_true",default = False, dest = "castor" , help = "get from castor")

    (options, args) = opt.parse_args()
    print args
    if options.print_usage :
        opt.print_help()
        exit()
    if len(args) != 1:
        opt.error("Incorrect number of arguments, should be one comma-separated list of ProdID(s).")

    # treating options
    global base_path, nb_logs, delta, MCversion
    MCversion = options.MCversion
    base_path =  os.path.abspath(os.curdir) #base_path = options.base_path
    nb_logs   = options.nb_logs 
    delta     = options.delta
    if not os.path.isabs( base_path ) :
        base_path = os.path.join( os.path.abspath(os.curdir), base_path )

    # treating args
    global list_Dirac_id
    if ',' not in args[0] : args[0] = args[0].replace(' ',',')
    list_Dirac_id = args[0].split(',')
    list_Dirac_id = map(int, list_Dirac_id )
   
    print " Dirac ids ", list_Dirac_id

    global castor_first 
    castor_first = options.castor

    


def take_info(file):
    """========================================================================
    Retrieve info from the .txt.gz or .log.bz2 logfile.
    ========================================================================"""
    
    import urllib, tempfile, os, gzip, bz2
    
    #print file
    log_file = ''
    
    tmpfile = tempfile.mktemp()
    
    try:
      urllib.urlretrieve(file,tmpfile)
    except IOError:
      print "Error: \tThe url %s is not available or doesn't exist.\n" % file
      return log_file
    
    # Open the log file. First try opening it as a gzip file and if that doesn't work try bz2.
    fc = None
    try:
      fc = gzip.open(tmpfile)
      log_file = fc.read()
    except IOError, e:
      # Close the gzip opened file and reopen it as bz2.
      fc.close()
      try:
        fc = bz2.BZ2File(tmpfile)
        log_file = fc.read()
      except IOError, e:
#        print "Error: %s could not be opened as a bz2 or gzip file. Unknown compression or file could not be found on server." %file
#        print e
# sm
        fc.close()
        try:
            fc = open(tmpfile)
            log_file = fc.read()
        except IOError, e:
            print "Error: %s could not be opened as uncompressed or a bz2 or gzip file. Unknown compression or file could not be found on server." %file
            print e
            
    if fc:
      fc.close()
      
    os.remove(tmpfile)
    return log_file



def getXMLMergeCommand():
    return os.path.join(base_path, 'XMLMerge.py')


#===========================================================
# 'MAIN' FUNCTION
#===========================================================    

parse_options()

print base_path , nb_logs, delta, list_Dirac_id


print '****************   create all folders *******************'

for prod in list_Dirac_id:
    path = os.path.join( base_path, repr(prod) )
    try: os.mkdir(path)
    except OSError:
        print 'File ', path, 'exists. Overwriting content.' 
        #try:  os.remove( os.path.join( base_path, repr(prod), "MergedSummary.xml" ) ) 
        #except OSError: 
        #    print "could not remove previous MergedSummary"
        #    continue
print '****************  download *******************'


#setup LHCbDirac
dirac_tmp_file_name = 'temp_SetupLHCbDirac'
# help from Hubert Degaudenzi
from LbConfiguration.SetupProject import SetupProject
r = SetupProject().main(    ['LHCbDirac'
                            , '--shell=csh'
                            ,'--output='+dirac_tmp_file_name
                            ])

envdict = os.environ.copy()
print len(envdict) 
df = open ( dirac_tmp_file_name , 'r' ) 
for line in df :
    if line.find("setenv") > -1 :
        s = line.split(" ") 
        #print s[0],s[1],s[2][1:-2]
        envdict[s[1]] = s[2][1:-2]

df.close()
print "after", len(envdict )

for prod in list_Dirac_id:
    
    #check if empty, inexistent. If so try to get them from castor
    # loop over the files and check if  404
    get_from_CASTOR = castor_first
    
    cmd_path = os.path.join( base_path, dirac_download_script )
    cwd_path = os.path.join( base_path, repr(prod) )
    
    if not get_from_CASTOR :
        
        print 'calling : python', cmd_path, str(prod), str(nb_logs), MCversion, "xml",  'from ', cwd_path 
        p = Popen(['python', cmd_path, str(prod), str(nb_logs), MCversion, "xml" ], stdout = PIPE, stderr = STDOUT,  cwd=cwd_path, env = envdict )#
        #printed = False
        print "printing lines of stdout"
	while True:
            line = p.stdout.readline()
            #if not printed :
            print line
            #    printed = True
            if not line: break
    
        # get the sizes of the logs
        sizes, files = [], []
        p = Popen(['ls', '-l'], stdout = PIPE, stderr = STDOUT, cwd=cwd_path )  # remove for ease of use
        broken = 0
        notbroken = 0
	while True:
            line = p.stdout.readline()
            if not line: 
		print "reached end of stdout"
		break
            l = line.split()
            print line,l,len(l)
            if len(l) == 9:
	    #if len(l) > 0:
                sizes.append(int(l[4]))
                files.append(l[8])
		notbroken += 1
            else: 
                broken += 1
        
        print "number of broken jobs : %d"%(broken)
	print "number of not-broken jobs = %d"%(notbroken)

    
        for f_name in files:
        
            try:
                f_test = open(os.path.join(cwd_path, f_name), 'r')
            except IOError:
                continue
            lines = f_test.readlines()
            if  '<title>404 Not Found</title>\n' in lines:
                f_test.close()
                os.remove(os.path.join(cwd_path, f_name))
                print 'the file does not exist, accessing castor'
                get_from_CASTOR = True
                break
    if get_from_CASTOR :
        #try to get files from CASTOR
        #try:
        print 'nsls', castor_location,'  | grep ', str(prod) 
        nsls_command = ['nsls', castor_location ]
        p = Popen( nsls_command, stdout = PIPE, stderr = STDOUT , env = envdict )
        lines = p.stdout.readlines()
        castor_files = []
        for l in lines:    
            if str(prod) in l:
                l=l.rstrip()
		l=l.translate(string.maketrans('',''),'/')
		castor_files.append(l)
		#castor_files.append( (l.rstrip()).translate(None,'/') )
        if len(castor_files) == 0:
            print 'This Prod ID is not archived on CASTOR either, sorry!'
        for castor_file in castor_files:
            #### look in current directories
            #get prod id form castor file
            local = False
            #print 'ls', os.path.join( ".", castor_file)
            #p = Popen(['ls', os.path.join( ".", castor_file) ], stdout = PIPE, stderr = STDOUT, cwd=cwd_path, stdin=PIPE  )
	    print 'nsls', os.path.join( castor_location, castor_file)
            p = Popen(['nsls', os.path.join( castor_location, castor_file) ], stdout = PIPE, stderr = STDOUT, cwd=cwd_path, stdin=PIPE  )
            line = p.stdout.readline()
            print local, line, p, line.find("No such")
            #if  line.find("No such") > -1 : #file or directory
            if not local :
                #need to copy it
                print 'rfcp '+castor_location+castor_file+' . in '+cwd_path
                #print 'rfcp '+castor_location+castor_file+' . in '+cwd_path
		p = Popen(['rfcp', os.path.join( castor_location, castor_file), '.'], stdout = PIPE, stderr = STDOUT, cwd=cwd_path, env = envdict )
                while True:
                    line = p.stdout.readline()
                    if not line: break
                    print line,
            else :
                print 'mv', os.path.join( cwd_path,"..", castor_file), os.path.join( cwd_path, castor_file ) 
                p = Popen(['mv', os.path.join( cwd_path,"..",  castor_file), os.path.join( cwd_path, castor_file ) ], stdout = PIPE, stderr = STDOUT, cwd=cwd_path)
                from time import sleep
                sleep(5)
        print "untarring", p
        count = 0
        for castor_file in castor_files:
            print "looking at castor file", castor_file
            ## copied from Vladimir Romanovsky's script Untar_Castor.py
            f = tarfile.open(os.path.join(cwd_path,castor_file),"r:gz")
            print "f = ", f
            names = f.getnames()
            for name in names:
		print "name = ", name
		######################################################################
		##	OLD DATA STRUCTURE - have found a tar file within the tar file
		######################################################################
                if name.count("tgz"):  
                    print "inside castor file", castor_file
                    print "find tar file", name
		    print "USING OLD DATA STRUCTURE e.g. MC11"
	 	    print "TARFILE ON CASTOR CONTAINS TARFILE, WHICH CONTAINS XML SUMMARY"
                    #name = os.path.split(name)[1]
		    #print "after split, name = ", name
		    #directory = os.path.dirname( os.path.join( cwd_path, castor_file ) )
		    #name = os.path.join( castor_location, name )
                    f.extract(name, path=cwd_path)
                    f2 = tarfile.open(os.path.join(cwd_path,name),"r:gz")
                    print "f2 = ", f2
                    name2s = f2.getnames()
                    for name2 in name2s:
			print "name2 = ", name2
			if name2.find("summaryDaVinci") == 0 and name2.count('xml'):
		            print "DV summary = ", name2
			    f2.extract(name2, path=cwd_path)
			    print "extracted  file ", name2
			    print "to directory ", cwd_path
			    if os.path.getsize(os.path.join(cwd_path,name2)) > 500:
			        count += 1
				print "count = ",count
				# move file to cwd directory - where log files are expected to be
				p = Popen(['mv', os.path.join( cwd_path,name), cwd_path ], stdout = PIPE, stderr = STDOUT, cwd=cwd_path)
			    else:  #ignore logs with too few events to be significant
			        try:
                                    os.remove(os.path.join(cwd_path,name))
				    print "removing file due to small number of evts: ",os.path.join(cwd_path,name)
                                except:
                                    pass
			else:
			    try:
			        os.remove(os.path.join(cwd_path,name))  # remove files that are not summaryDaVinci.xml
			    except:
			        pass
			    try:
			        os.rmdir(os.path.split(os.path.join(cwd_path,name))[0]) # remove directory after all logs within have been read
			    except:
			        pass
		########################################################################
		##  		NEW DATA STRUCTURE
		#########################################################################
		else:
                    f2 = f
                    name2 = name 
                if name2.find("summaryDaVinci") == 23 and name2.count('xml'):
                    print "USING NEW DATA STRUCTURE e.g. MC12"
                    print "TARFILE ON CASTOR DIRECTLY CONTAINS XML SUMMARY"
                    print "name2 = ", name2
                    f2.extract(name2, path=cwd_path)
                    print "extracted  file ", name2
                    print "to directory ", cwd_path
                    if os.path.getsize(os.path.join(cwd_path,name)) > 500:
                        count += 1
                        print "count = ",count
                        # move file to cwd directory - where log files are expected to be
                        p = Popen(['mv', os.path.join( cwd_path,name), cwd_path ], stdout = PIPE, stderr = STDOUT, cwd=cwd_path)
                    else:  #ignore logs with too few events to be significant
                        try:
                            os.remove(os.path.join(cwd_path,name))
                            print "removing file due to small number of evts: ",os.path.join(cwd_path,name) 
                        except:
                            pass 
                else:
                    try:
                        os.remove(os.path.join(cwd_path,name))  # remove files that are not summaryDaVinci.xml
                    except:
                        pass
                    try:
                        os.rmdir(os.path.split(os.path.join(cwd_path,name))[0]) # remove directory after all logs within have been read
                    except:
                        pass
                if count > nb_logs: 
                    print "count > nb_logs, break"
                    break
            #don't remove original file         
            if False: #not local :
                try:
                    os.remove(os.path.join(cwd_path,castor_file))
                except:
                    pass
            ##else :
            ##    print 'mv', os.path.join( cwd_path , castor_file), os.path.join( "." , castor_file ) 
            ##    p = Popen(['mv', os.path.join( cwd_path , castor_file), os.path.join( "." , castor_file ) ], stdout = PIPE, stderr = STDOUT, cwd=cwd_path)


            #shutil.rmtree()
            ## end of Vladimir Romanovsky's script Untar_Castor.py
            #import code ; code.interact( local=locals())

        print "get the sizes of the logs"
        sizes, files = [], []
        p = Popen(['ls', '-l'], stdout = PIPE, stderr = STDOUT, cwd=cwd_path)
        while True:
            line = p.stdout.readline()
            if not line: break
            #print line
            if line.find("tgz") > -1 :
                continue
            l = line.split()
            if len(l) == 9:
                #print line,l, l[4], l[8]
                sizes.append(int(l[4]))
                files.append(l[8])
            else: 
				pass
            
    print max(sizes), delta*max(sizes)
    #sys.exit(2) 
    print '-- cleaning up logs.'
    
    for i in range(len(sizes)):
        #rm logs with sizes less than max-0.07*max
        ##if sizes[i] > 400 : #int(max(sizes)-delta*max(sizes)):
            #print 'removing', os.path.join(cwd_path, files[i]), 'because it is', sizes[i], 'large; minimal size is', int(max(sizes)-delta*max(sizes))
        ##    if files[i].find(".tgz") > -1 :
        ##        print "do not remove", files[i]
        ##        continue
        ##    try:
        ##        os.remove( os.path.join(cwd_path, files[i]))
        ##        print "removing file ", os.path.join(cwd_path, files[i])
        ##    except OSError:
        ##        print "OSError removing",os.path.join(cwd_path, files[i])
                #pass
        ##        try:
        ##            shutil.rmtree( os.path.join(cwd_path, files[i]), True)            
        ##        except OSError :
        ##            print "OSError, can't remove", os.path.join(cwd_path, files[i])
            

        #rm files with size == 358 (empty logs)
        ##elif sizes[i] < 400:
        if sizes[i] < 400:
            print 'removing', os.path.join(cwd_path, files[i]), 'because it is empty, its size is', sizes[i]
            os.remove( os.path.join(cwd_path, files[i]) )

    # check if there any logs left
    if len(os.listdir(cwd_path)) == 0:
        print '\nNo logs left in', cwd_path
        print 'Will ignore Prod ID', prod
        list_Dirac_id.remove(prod)
        #sys.exit()
    
print '****************  run XMLMerge.py *******************'

XMLMerge_cmd = getXMLMergeCommand()


#setup DaVinci
davinci_tmp_file_name = 'temp_SetupDaVinci'
# help from Hubert Degaudenzi
from LbConfiguration.SetupProject import SetupProject
r = SetupProject().main(    ['DaVinci'
                            , '--shell=csh'
                            ,'--output='+davinci_tmp_file_name
                            ])

envdict = os.environ.copy()
print "before", len(envdict)
df = open ( davinci_tmp_file_name , 'r' ) 
for line in df :
    if line.find("setenv") > -1 :
        s = line.split(" ") 
        envdict[s[1]] = s[2][1:-2]
df.close()
print "after", len(envdict)


# do the merging

for prod in list_Dirac_id:

    files = []
    for file in os.listdir(os.path.join( base_path, repr(prod))):
        if file[-4:] != '.xml':
            continue
        #print file, os.path.join( base_path, repr(prod), file)
        files.append(os.path.join(file))
    if len(set(files)) ==0 :
        print 'No xml files, skipping Prod ID', prod
        continue 
   
    #print files
    cwd_path = os.path.join( base_path, repr(prod) )
    cmdlist = ["python", XMLMerge_cmd ] + files
    cmdline = "" 
    for cmd in cmdlist : cmdline += cmd + " "
    print "running", cmdline

    p = Popen( cmdlist , stdout = PIPE, stderr = STDOUT,  cwd=cwd_path, env = envdict )
    while True:
        line = p.stdout.readline()
        if not line : break 
        print line [-1]



try:
    os.remove( dirac_tmp_file_name )
except:
    pass

try:
    os.remove( davinci_tmp_file_name )
except:
    pass


for d in os.listdir(".") :
    if os.path.isdir(d) :
        for f in os.listdir(d) :
            if f.find("Merged") > -1  :
                cmd = "cp "+os.path.join(d,f)+" "+f[ : f.find(".xml") ] + "_" + d + ".xml" 
                print cmd
                os.system(cmd) 

print "cwd_path = ", cwd_path
        
