#! /usr/bin/env python2
# This file is released under the BSD licence.

__author__ = 'Greig A Cowan, Sam Skipsey, Wahid Bhimji'
__date__ = 'September 2009'
__version__ = 0.1

'''Look for files in a particular spacetoken and replicate'''

import sys
import os
import string
import getpass
import MySQLdb
import optparse
import gridpp_dpm

def main():

    usage = '''usage: %prog --st SPACETOKEN --nreps N(=2)

    This allows you to replicate files in a given spacetoken'''

    # Add check_required to option parser
    class OptionParser (optparse.OptionParser):
        def check_required (self, opt):
            option = self.get_option(opt)            
            # Assumes the option's 'default' is set to None!
            if getattr(self.values, option.dest) is None:
                self.error("%s option required and not supplied. Try --help." % option)
                
    parser = OptionParser( usage = usage)
    parser.add_option('--st', dest='st',default=None,
                      help='Specify a space token description')
    parser.add_option('--nreps', dest='nreps',default=2,
                      help='Specify the number of copies required.Default 2.')
    parser.add_option('--del', action="store_true", dest="dodel", default=False,
                      help='Delete excess replicas (above amount specified in nreps)')
    parser.add_option('--list', action="store_true", dest="dolist", default=False,
                      help='Just list replicas. No action taken.')
    parser.add_option('--verbose', action="store_true", dest="doverbose", default=False,
                      help='Print more output')
   
    (options, args) = parser.parse_args()
    parser.check_required("--st")
    
    if len(args) > 0:
        parser.error("incorrect number of arguments")
        exit(1)

    (c,cc,dpmdb) = gridpp_dpm.MySQLConnect(True)


    # Read DB to find list of files in space token and number of replicas
        
    try:
        cc.execute('''
         select GROUP_CONCAT(Cns_file_replica.sfn), Cns_file_replica.setname, count(*) from Cns_file_replica join %(dpm_db)s.dpm_space_reserv on %(dpm_db)s.dpm_space_reserv.s_token = Cns_file_replica.setname where %(dpm_db)s.dpm_space_reserv.u_token="%(st)s" group by Cns_file_replica.fileid''' % {"st":options.st,"dpm_db":dpmdb})

        results = cc.fetchall()
    except MySQLdb.Error, e:
        print "Error %d: %s" % (e.args[0], e.args[1])
        sys.exit (1)
    cc.close()
    c.close()

    #Use dpm-replicate as addreplica doesn't take spacetoken

    try:
        dpmrepcmd = os.environ['LCG_LOCATION'].rstrip('/') + '/bin/dpm-replicate --f_type PERMANENT --space_token'
    except:
        dpmrepcmd = '/opt/lcg/bin/dpm-replicate --f_type PERMANENT --space_token'
    if (options.dodel == True):
        try:
            dpmdelcmd = os.environ['LCG_LOCATION'].rstrip('/') + '/bin/dpm-delreplica'
        except:
            dpmdelcmd = '/opt/lcg/bin/dpm-delreplica'
            
    for line in results:
	try: #Python-MySQL 1.2.1 workaround, as this returns c arrays for blobs, not strings
        	repfns = line[0].tostring().split(',')
	except:
		repfns = line[0].split(',')
        repfn = repfns[0]
        reps2 = len(repfns)
        stoken = line[1]
        reps = line[2]
        if (options.dolist == True):
            print "%s replicas for file which are: " % reps
            print repfns
            continue 
        if int (reps2) == int (reps) :
            if (options.dodel == True):
                while (int(options.nreps) - int(reps)) < 0:
                    print  "%s replica(s) for %s so deleting" % (reps,repfns[reps-1])
                    os.system('%s %s' % (dpmdelcmd,repfns[reps-1]))
                    reps = reps - 1
            else:         
                while (int(options.nreps) - int(reps)) > 0 :
                    print "%s replica(s) for %s so replicating" % (reps,repfn)
                    os.system('%s %s %s' % (dpmrepcmd,stoken,repfn))
                    reps = reps + 1
            if  (options.doverbose == True):
                print "Enough reps for %s" % repfn
        else:
            print "Number of replicas listed does not match count in Db query. Skipping " 
            print line
            
if __name__ == '__main__':
         main()
