#!/usr/bin/env python
"""
MigrateToSWTEST: a script to copy/replicate files in the TestFileDB from
                 wherever they are into SWTEST

usage: MigrateToSWTEST.py [testname] [--help] [--verbose] [--dryrun] [--no-guid]

You require an environment with LHCbDirac and some package which depends on PRCOnfig to run this script, and a local copy of PRConfig.

If --no-guid is given, the file will be uploaded with a random GUID, not obtained from the file itself

e.g. SetupProject LHCbDirac
     lhcb-proxy-init [--group lhcb_prod]
     SetupProject LHCb ... --build-env --use PRConfig --use AppConfig
     getpack PRConfig head
     SetupProject LHCb ... --use PRConfig --use AppConfig

where "LHCb" can be any LHCb project or application.

if no testname is given, you will be propted to choose one. Currently only local files can be migrated, ROOT files will have their GUID extracted.

GRID files not on EOS/SWTEST will be checked to see first if they are on SWTest

See the Twiki: https://twiki.cern.ch/twiki/bin/view/LHCb/TestFileDB
"""
import sys
import os
from time import sleep

import commands, tempfile

DRYRUN=False
isVerbose=False
withGUID=True

enableReplicate=False

def usage():
    print __doc__

if "--help" in sys.argv or "-h" in sys.argv:
    usage()
    sys.exit(0)
if "--verbose" in sys.argv or "-v" in sys.argv:
    isVerbose=True
if "--dryrun" in sys.argv or "-d" in sys.argv:
    DRYRUN=True
if "--no-guid" in sys.argv or "-g" in sys.argv:
    withGUID=False


######### Check environment ###########

try:
    import PRConfig
    from PRConfig import TestFileDB
    from PRConfig.TestFileDB import test_file_db
except ImportError:
    print usage()
    raise ImportError("No PRConfig package could be loaded, your environment is incorrect")

if not os.access(PRConfig.__file__, os.W_OK):
    usage()
    raise IOError("You do not have permission to edit the PRConfig database, probably you're not working with a local PRConfig version. The version you are picking up is in "+PRConfig.__file__)

if "LHCBDIRACSYSROOT" not in os.environ:
    usage()
    raise AttributeError("Doesn't look to me like you have Setup the LHCbDirac environment correctly, please try again!")

if "APPCONFIGROOT" not in os.environ:
    usage()
    raise AttributeError("Currently I find the file GUID running a script from AppConfig, but you don't have AppConfig in your environment, please add --use AppConfig")

try:
    import GaudiConf
except ImportError:
    print usage()
    raise ImportError("Doesn't look like you have IOHelper available, so you're probably not running in an LHCb project, please try again with the corect environment")

try:
    import DIRAC
except ImportError:
    raise ImportError("Doesn't look to me like you have Setup the LHCbDirac environment correctly, please try again!")



#check for proxy
stat,out=commands.getstatusoutput("lhcb-proxy-info")
if isVerbose or DRYRUN:
    print "-----------------"
    print stat, out
    print "-----------------"

if "Proxy is expired" in out:
    usage()
    raise AttributeError("Your grid proxy is expired, pleae create it first")

if "lhcb_user" in out:
    print "found lhcb-proxy, group is User, this means you will not be able to replicate production files, but you will be able to upload user files"
    
elif "lhcb_prod" in out:
    print "lhcb_prod recognised, replication of production files enabled, but you will not be able to replicate user files"
    enableReplicate=True

#########   Steering, patterns to identify different file types #########

tGrid=[]; nGrid=0; grid="/lhcb/grid/";
tOnEOS=[]; nOnEOS=0; eos="root://eos"
tOnCastor=[]; nOnCastor=0; castor="/castor/cern.ch/"
tOnSWTEST=[]; nOnSWTEST=0; swtest="/lhcb/swtest/"
tOnAfs=[]; nOnAfs=0; afs="/afs/cern.ch/"
tOnLocal=[]; nOnLocal=0
nTotal=0

for k,v in test_file_db.iteritems():
    nTotal=nTotal+len(v.filenames)
    flag=False
    if grid in v.filenames[0]:
        tGrid.append(k)
        #print v.filenames
        nGrid=nGrid+len(v.filenames)
        flag=True
    if eos in v.filenames[0]:
        tOnEOS.append(k)
        nOnEOS=nOnEOS+len(v.filenames)
        flag=True
    if swtest in v.filenames[0]:
        tOnSWTEST.append(k)
        nOnSWTEST=nOnSWTEST+len(v.filenames)
        flag=True
    if castor in v.filenames[0]:
        tOnCastor.append(k)
        nOnCastor=nOnCastor+len(v.filenames)
        flag=True
    if afs in v.filenames[0]:
        tOnAfs.append(k)
        nOnAfs=nOnAfs+len(v.filenames)
        flag=True
    if not flag:
        tOnLocal.append(k)
        nOnLocal=nOnLocal+len(v.filenames)

print "=============== SUMMARY ================="
print "Inclusive list of file locations"
print "Location,  #of testfile objects,  #of files"
print "Grid  : ", "\t\t", len(tGrid),"\t\t", nGrid
print "EOS   : ", "\t\t", len(tOnEOS),"\t\t", nOnEOS
print "SWTEST: ","\t\t", len(tOnSWTEST),"\t\t", nOnSWTEST
print "Castor: ","\t\t", len(tOnCastor),"\t\t", nOnCastor
print "Afs   : ","\t\t", len(tOnAfs),"\t\t", nOnAfs
print "Local : ","\t\t", len(tOnLocal),"\t\t", nOnLocal
print "----------------------------------------"
print "Total :","\t\t", len(test_file_db),"\t\t", nTotal
print "=============== TO MIGRATE ============="
print "Castor:",
if not len(tOnCastor):
    print "None"
else:
    print
for t in tOnCastor:
    print "  -->", t
print "Afs:",
if not len(tOnAfs):
    print "None"
else:
    print
for t in tOnAfs:
    print "  -->", t
print "Local:",
if not len(tOnLocal):
    print "None"
else:
    print
for t in tOnLocal:
    print "  -->", t
replicate=[t for t in tGrid if t not in tOnSWTEST]
if enableReplicate:
    print "On Grid, but not SWTest:",
    if not len(replicate):
        print "None"
    else:
        print
    for t in replicate:
        print "  -->", t

print "=============== MIGRATION ================="

if len(tOnAfs+tOnCastor+tOnLocal+replicate)<2:
    print "no suitable tests for migration. Add first using the AddToTestFilesDB.py script"
    sys.exit(0)

testtomigrate=None
if len(sys.argv)>1 and len([s for s in sys.argv[1:] if s not in ["-v","-d","--verbose","--dryrun"]]):
    testtomigrate=[s for s in sys.argv[1:] if s not in ["-v","-d","--verbose","--dryrun"]][0]
else:
    testtomigrate=raw_input("Choose testfile name to migrate: ")

if testtomigrate not in test_file_db.keys():
    usage()
    raise ValueError("That test does not exist "+testtomigrate)

if (not enableReplicate) and (testtomigrate not in tOnCastor+tOnAfs+tOnLocal or testtomigrate in tGrid):
    usage()
    raise NameError("Currently can only migrate from castor/afs/local to EOS-Grid-SWTest, because you only have the lhcb_user group, either remake your proxy with --group lhcb_prod, or contact Philippe Charpentier to address this issue.")
elif enableReplicate and (testtomigrate not in tOnCastor+tOnAfs+tOnLocal+replicate):
    usage()
    raise NameError("Test "+testtomigrate+" does not require migration")

print "Will relocate", len(test_file_db[testtomigrate].filenames), "file of the form", test_file_db[testtomigrate].filenames[0], "to EOS-GRID-SWTEST"
yn=raw_input("Is this OK? y/[n]").strip().lower()
if not yn.startswith("y"):
    sys.exit(0)

###########   User has chosen migration :) Let's get on with it... #####

from GaudiConf import IOHelper


def prepare(f,tmpdir=None):
    "prepare file for transfer, copy to temp if necessary, split off any protocols"
    realf=IOHelper().undressFile(f)
    realf=f.split(':')[-1].split('?')[0].replace("//castorlhcb.cern.ch//","/").replace("//castor.cern.ch//","/")
    if castor in f and tmpdir is not None:
        if isVerbose:
            print "----------------------------"
        print "Copying to temp..."
        cmd="rfcp "+realf+" "+tmpdir+"/"
        if DRYRUN or isVerbose:
            print cmd
        if not DRYRUN:
            os.system(cmd)
        return tmpdir+"/"+realf.split('/')[-1]
    return realf

def constructlfn(f,testname):
    """
    Automate creating a new lfn for uploading to swtest
    /lhcb/swtest/testname/filename
    """
    return "/lhcb/swtest/"+testname+"/"+f.split('/')[-1].split("?")[0]

def readGUID(f):
    """
    Use the AppConfig script to determine the gaudi GUID
    """
    cmd=os.environ["APPCONFIGROOT"]+"/scripts/ReadGUID.py "+f.replace("PFN:","")
    if isVerbose:
        print "----------------------------"
        print "Trying to read the GUID using AppConfig script"
        print cmd
    stat,out=commands.getstatusoutput(cmd)
    if stat or isVerbose:
        print stat, out
    if stat:
        raise OSError("couldn't run AppConfig script to get the file ID")
    return out.strip().encode('ascii').replace("\0","").replace("[?1034h","").strip().split(".")[0]

def copy(f,lfn,guid=None):
    """
    Actually call the dirac-dms command
    """
    if isVerbose:
        print "----------------------------"        
    print "Uploading...",
    if guid is not None:
        print guid,
    print
    acommand="dirac-dms-add-file "+lfn+" "+f+" CERN-SWTEST"
    if guid is not None:
        acommand=acommand+" "+guid
    acommand=acommand[:]
    if DRYRUN or isVerbose:
        print acommand
        for s in acommand:
            print s,
        print
    if not DRYRUN:
        try:
            stat,out=commands.getstatusoutput(acommand)
        except:
            atexit()
            print "trying", acommand
            raise
        if isVerbose:
            print stat
        if "GUID already exists " in out:
            line=[l for l in out.split('\n') if "GUID already exists " in l]
            lfn="/"+"/".join(line[0].split("/")[1:])
            lfn=lfn.strip()
            if not enableReplicate:
                atexit()
                raise IOError("File already exists in the book-keeping, or at least a file with this same GUID. You don't have the correct Dirac group to perform replication, either start again with lhcb_prod or contact Philippe Charpentier to replicate this file. "+guid+" -> "+lfn)
            else:
                print "GUID already exists on the grid as : ",lfn
                replicate(lfn)
        elif "Write access not permitted for this credential" in out:
            atexit()
            print out
            raise IOError("There was a permissions error on trying to upload the file, probably you don't have the correct dirac role. If you are trying to upload a user file, you may want to try again with a user proxy.")
        else:
            print out
        sleep(5) #registration time delay
        #sleep(5)
    return lfn

def checkongrid(pfn):
    """
    double check if this file is actually on the grid already, perhaps I can't/didn't guess the lfn corectly
    """
    cmd="dirac-dms-lfn-replicas "+pfn
    if isVerbose:
        print "----------------------------"
        print "Checking file pfn is on grid"
        print cmd
    stat,out=commands.getstatusoutput(cmd)
    if isVerbose:
        print stat, out
    lfn=False
    if "Successful" in out:
        lfn=out.split(":")[1].strip()
    return lfn

def check(lfn):
    """
    Is this LFN on CERN-SWTEST
    """
    cmd="dirac-dms-lfn-replicas "+lfn
    if isVerbose:
        print "----------------------------"
        print "Checking file lfn is in SWTEST "
        print cmd
    
    stat,out=commands.getstatusoutput(cmd)
    if isVerbose:
        print stat, out
    return ("CERN-SWTEST" in out)

def geturl(lfn):
    """
    Dermine the accessURL for this LFN at CERN-SWTEST
    """
    cmd="dirac-dms-lfn-accessURL "+lfn+" CERN-SWTEST"
    if isVerbose:
        print "----------------------------"
        print "determining access URL"
    if DRYRUN or isVerbose:
        print cmd
    stat,out=commands.getstatusoutput(cmd)
    if DRYRUN or isVerbose:
        print stat, out
    if DRYRUN:
        if "Error " in out:
            out=lfn+" : not/a/pfn/dry/run"
    elif stat!=0 or "Error " in out:
        raise IOError(out)
    url=out.split('\n')[-1].replace(lfn+" : ","").strip()
    if (lfn.endswith('.raw') or lfn.endswith('.mdf')) and url.startswith("root:"):
        #temporarily fix Dirac problem
        url="mdf:"+url
    return url

def replicate(lfn):
    """
    replicate an existing grid file into swtest
    """
    
    cmd="dirac-dms-replicate-lfn "+lfn+" CERN-SWTEST"
    if isVerbose:
        print "----------------------------"
    print "Replicating ..."
    if DRYRUN or isVerbose:
        print cmd
    stat,out=(0,"TEST")
    if not DRYRUN:
        stat,out=commands.getstatusoutput(cmd)
    if DRYRUN or isVerbose:
        print stat, out
    return lfn

def mktmpdir():
    """
    make a temporary file location
    """
    return tempfile.mkdtemp()


def pintodb(replacements):
    """
    Compy this modification into the database
    """
    import PRConfig
    yn=raw_input("Pin to the database? y/[n]").strip().lower()
    if not yn.startswith("y"):
        return False
    dbname="/".join(PRConfig.__file__.split('/')[:-1])+"/TestFileDB.py"
    f=open(dbname)
    original=f.read()
    f.close()
    update=original
    for rep in replacements:
        update=update.replace(rep,replacements[rep])
    if not DRYRUN:
        f=open(dbname,"w")
        f.write(update)
        f.close()
    if DRYRUN or isVerbose:
        print update
    return True

def checkdb(testtomigrate,oldtest):
    oldfiles=test_file_db[testtomigrate].filenames[:]
    import PRConfig
    from PRConfig import TestFileDB
    reload(PRConfig)
    reload(TestFileDB)
    if isVerbose or DRYRUN or oldfiles!=TestFileDB.test_file_db[testtomigrate].filenames:
        print "New test entry appears as"
        print TestFileDB.test_file_db[testtomigrate].recreateme()
    if oldfiles!=TestFileDB.test_file_db[testtomigrate].filenames and not DRYRUN:
        atexit()
        raise NameError("Pinning did not work, probably this means there is some logic in the python of the database, which you need to update by hand")
    elif oldfiles!=TestFileDB.test_file_db[testtomigrate].filenames:
        print "Pinning did not work, probably this means there is some logic in the python of the database, which you need to update by hand"
    return True


atmp=None

if testtomigrate in tOnCastor:
    atmp=mktmpdir()
    #should now delete this directory, no matter what happens!

#cleanup temp
def atexit():
    """
    Cleanup tmp files at exit in case something else goes wrong, don't want to fill up the disk!
    """
    if isVerbose:
        print "Exit handler invoked, cleaning"
    if atmp:
        if os.path.exists(atmp):
            if len(atmp):
                if atmp.startswith("/tmp"):
                    os.system("rm -rf "+atmp)



replacements={}

files=test_file_db[testtomigrate].filenames
if DRYRUN:
    files=files[:1]

#detect the file format automatically
from GaudiConf import IOExtension
aformat=IOExtension().detectMinType(files)

for f in files:
    if isVerbose:
        print "----------------------------"
        print "Migrating "+f
    lfn=checkongrid(f)
    ######## Branch if file is already on the grid
    if lfn is not False:
        print "PFN is already on the grid, at lfn:", lfn
        if not enableReplicate:
            usage()
            atexit()
            raise ValueError("you should not be trying to replicate a file because you haven't got the correct group, try lhcb-proxy-init --group lhcb_prod or contact Phillipe Charpentier")
        if isVerbose:
            print "Migrating to "+lfn
        if check(lfn):
            print "LFN already exists, migrating database only"
        else:
            if isVerbose:
                print "Preparing to replicate "+f+" -> CERN-SWTEST"
            replicate(lfn)
    
    else:
        #is the file already on the grid?
        #if not, make the lfn myself
        lfn=constructlfn(f,testtomigrate)
        if isVerbose:
            print "Migrating to "+lfn
        
        if check(lfn):
            print "LFN already exists, migrating database only"
        else:
            if isVerbose:
                print "Preparing to upload "+f+" -> "+lfn 
            GUID=None
            if aformat=="ROOT" and withGUID:
                GUID=readGUID(f)
            realf=prepare(f,atmp)
            lfn=copy(realf,lfn,GUID)
            if atmp is not None:
                if len(atmp)>2:
                    if atmp in realf:
                        if atmp.startswith('tmp'):
                            os.system("rm "+realf)
    c=check(lfn)
    if not c and not DRYRUN:
        atexit()
        raise IOError("Unable to migrate file, was not successful "+f+" -> "+lfn)
    elif not c:
        print "Unable to migrate file, was not successful "+f+" -> "+lfn
    replacements[f]=geturl(lfn)

print "=============== FINALIZING ================="
print "Copying successful, new testfile object will have modified files"

newnames=[]
for k,v in replacements.iteritems():
    newnames.append(v)

test_file_db[testtomigrate].filenames=newnames

if isVerbose or DRYRUN:
    print test_file_db[testtomigrate].recreateme()

if pintodb(replacements):
    if checkdb(testtomigrate,test_file_db[testtomigrate]):
        print "Done. You can check the files are readable using $PRCONFIGROOT/scripts/TestReadable.py"

atexit()
