
# Dirac commands
#/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/

def getJobGroupJobs(jg):
    ''' Return jobs in a group'''
    output(dirac.selectJobs(jobGroup=jg))


def kill(id):
    ''' Kill a given DIRAC Job ID within DIRAC '''
    output(dirac.delete(id))


def peek(id):
    ''' Peek at the DIRAC Job id and return what we saw '''
    output(dirac.peek(id))


def getJobCPUTime(id):
    ''' Get the amount of CPU time taken by the DIRAC Job id'''
    output(dirac.getJobCPUTime(id))


def reschedule(id):
    ''' Reschedule within DIRAC a given DIRAC Job id'''
    output(dirac.reschedule(id))


def submit(djob, mode='wms'):
    ''' Submit a DIRAC job given by the jdl:djob with a given mode '''
    output(dirac.submit(djob, mode=mode))


def ping(system, service):
    ''' Ping a given service on a given system running DIRAC '''
    output(dirac.ping(system, service))


def removeFile(lfn):
    ''' Remove a given LFN from the DFC'''
    ret = {}
    if type(lfn) is list:
        for l in lfn:
            ret.update(dirac.removeFile(l))
    else:
        ret.update(dirac.removeFile(lfn))
    output(ret)


def getMetadata(lfn):
    ''' Return the metadata associated with a given :DN'''
    output(dirac.getMetadata(lfn))


def getReplicas(lfns):
    ''' Return  the locations of the replicas of a given LFN in a dict format, SE: location '''
    output(dirac.getReplicas(lfns, active=True, preferDisk = True))


def getAccessURL(lfn,SE):
    ''' Return the access URL for the given LFN and storage element '''
    output(dirac.getAccessURL(lfn,SE))


def getFile(lfns, destDir=''):
    ''' Put the physical file behind the LFN in the destDir path'''
    output(dirac.getFile(lfns, destDir=destDir))


def replicateFile(lfn, destSE, srcSE='', locCache=''):
    ''' Replicate a given LFN from a srcSE to a destSE'''
    res = dirac.replicateFile(lfn, destSE, srcSE, locCache)
    output(res)
    #print(res)


def removeReplica(lfn, sE):
    ''' Remove the physical files and LFN from the DFC'''
    output(dirac.removeReplica(lfn, sE))


def getOutputData(id, outputFiles='', destinationDir=''):
    ''' Return output data of a requeted DIRAC Job id, place outputFiles in a given destinationDir') '''
    output(dirac.getJobOutputData(id, outputFiles, destinationDir))


def splitInputData(files, files_per_job):
    ''' Split list of files ito a list of list of smaller files (below files_per_job in length) and return the list of lists'''
    output(dirac.splitInputData(files, files_per_job))


def getInputDataCatalog(lfns, site, xml_file):
    ''' Get the XML describing the given LFNs at a given site'''
    output(dirac.getInputDataCatalog(lfns, site, xml_file))


def uploadFile(lfn, file, diracSEs, guid=None):
    ''' Upload a given file to an lfn with 1 replica places at each element in diracSEs. Use a given guid if given'''
    outerr = {}
    for se in diracSEs:
        result = dirac.addFile(lfn, file, se, guid)
        if result.get('OK', False) and lfn in result.get('Value', {'Successful': {}})['Successful']:
            result['Value']['Successful'][lfn].update({'DiracSE': se})
            md = dirac.getMetadata(lfn)
            if md.get('OK', False) and lfn in md.get('Value', {'Successful': {}})['Successful']:
                guid = md['Value']['Successful'][lfn]['GUID']
                result['Value']['Successful'][lfn].update({'GUID': guid})
            output(result)
            return
        outerr.update({se: result})
    else:
        output(outerr)


def addFile(lfn, file, diracSE, guid):
    ''' Upload a given file to an lfn with 1 replica places at each element in diracSEs. Use a given guid if given'''
    output(dirac.addFile(lfn, file, diracSE, guid))


def getOutputSandbox(id, outputDir=os.getcwd(), oversized=True, noJobDir=True, pipe_out=True):
    '''
    Get the outputsandbox and return the output from Dirac to the calling function
    id: the DIRAC jobid of interest
    outputDir: output directory locall on disk to use
    oversized: is this output sandbox oversized this will be modified
    noJobDir: should we create a folder with the DIRAC job ID?
    output: should I output the Dirac output or should I return a python object (False)'''
    result = dirac.getOutputSandbox(id, outputDir, oversized, noJobDir)
    if result is not None and result.get('OK', False):

        if not noJobDir:
            tmpdir = os.path.join(outputDir, str(id))
            os.system('mv -f %s/* %s/. ; rm -rf %s' % (tmpdir, outputDir, tmpdir))
        
        os.system('for file in $(ls %s/*_Ganga_*.log); do ln -s ${file} %s/stdout; break; done' % (outputDir, outputDir))

    if pipe_out:
        output(result)
    else:
        return result


def getOutputDataInfo(id, pipe_out=True):
    ''' Get information on the output data generated by a job of ID and pipe it out or return it'''
    ret = {}
    result = getOutputDataLFNs(id, pipe_out=False)
    if result.get('OK', False) and 'Value' in result:
        for lfn in result.get('Value', []):
            file_name = os.path.basename(lfn)
            ret.update({file_name: {'LFN': lfn}})
            md = dirac.getMetadata(lfn)
            if md.get('OK', False) and lfn in md.get('Value', {'Successful': {}})['Successful']:
                ret[file_name].update(
                    {'GUID': md['Value']['Successful'][lfn]['GUID']})
            # this catches if fail upload, note lfn still exists in list as
            # dirac tried it
            elif md.get('OK', False) and lfn in md.get('Value', {'Failed': {}})['Failed']:
                ret[file_name].update({'LFN': '###FAILED###'})
                ret[file_name].update(
                    {'LOCATIONS': md['Value']['Failed'][lfn]})
                ret[file_name].update({'GUID': 'NotAvailable'})
                continue
            rp = dirac.getReplicas(lfn)
            if rp.get('OK', False) and lfn in rp.get('Value', {'Successful': {}})['Successful']:
                ret[file_name].update(
                    {'LOCATIONS': rp['Value']['Successful'][lfn].keys()})
    if pipe_out:
        output(ret)
    else:
        return ret


# could shrink this with dirac.getJobOutputLFNs from ##dirac
def getOutputDataLFNs(id, pipe_out=True):
    ''' Get the outputDataLFN which have been generated by a Dirac job of ID and pipe it out or return it'''
    parameters = dirac.parameters(id)
    lfns = []
    ok = False
    message = 'The outputdata LFNs could not be found.'

    if parameters is not None and parameters.get('OK', False):
        parameters = parameters['Value']
        # remove the sandbox if it has been uploaded
        sandbox = None
        if 'OutputSandboxLFN' in parameters:
            sandbox = parameters['OutputSandboxLFN']

        # now find out about the outputdata
        if 'UploadedOutputData' in parameters:
            lfn_list = parameters['UploadedOutputData']
            import re
            lfns = re.split(',\s*', lfn_list)
            if sandbox is not None and sandbox in lfns:
                lfns.remove(sandbox)
            ok = True
        elif parameters is not None and 'Message' in parameters:
            message = parameters['Message']

    result = {'OK': ok}
    if ok:
        result['Value'] = lfns
    else:
        result['Message'] = message

    if pipe_out:
        output(result)
    else:
        return result


def normCPUTime(id, pipe_out=True):
    ''' Get the normalied CPU time that has been used by a DIRAC job of ID and pipe it out or return it'''
    parameters = dirac.parameters(id)
    ncput = None
    if parameters is not None and parameters.get('OK', False):
        parameters = parameters['Value']
        if 'NormCPUTime(s)' in parameters:
            ncput = parameters['NormCPUTime(s)']
    if pipe_out:
        output(ncput)
    else:
        return ncput


def finished_job(id, outputDir=os.getcwd(), oversized=True, noJobDir=True):
    ''' Nesting function to reduce number of calls made against DIRAC when finalising a job, takes arguments such as getOutputSandbox
    Returns the CPU time of the job as a dict, the output sandbox information in another dict and a dict of the LFN of any uploaded data'''
    out_cpuTime = normCPUTime(id, pipe_out=False)
    out_sandbox = getOutputSandbox(id, outputDir, oversized, noJobDir, pipe_out=False)
    out_dataInfo = getOutputDataInfo(id, pipe_out=False)
    outStateTime = {'completed' : getStateTime(id, 'completed', pipe_out=False)}
    output((out_cpuTime, out_sandbox, out_dataInfo, outStateTime))


def status(job_ids, statusmapping, pipe_out=True):
    '''Function to check the statuses and return the Ganga status of a job after looking it's DIRAC status against a Ganga one'''
    # Translate between the many statuses in DIRAC and the few in Ganga

    result = dirac.status(job_ids)
    if not result['OK']:
        if pipe_out:
            output(result)
            return
        else:
            return result
    status_list = []
    bulk_status = result['Value']
    for _id in job_ids:
        job_status = bulk_status.get(_id, {})
        minor_status = job_status.get('MinorStatus', None)
        dirac_status = job_status.get('Status', None)
        dirac_site = job_status.get('Site', None)
        ganga_status = statusmapping.get(dirac_status, None)
        if ganga_status is None:
            ganga_status = 'failed'
            dirac_status = 'Unknown: No status for Job'
        #if dirac_status == 'Completed' and (minor_status not in ['Pending Requests']):
        #    ganga_status = 'running'
        if minor_status in ['Uploading Output Data']:
            ganga_status = 'running'

        try:
            from DIRAC.Core.DISET.RPCClient import RPCClient
            monitoring = RPCClient('WorkloadManagement/JobMonitoring')
            app_status = monitoring.getJobAttributes(_id)['Value']['ApplicationStatus']
        except:
            app_status = "unknown ApplicationStatus"

        status_list.append([minor_status, dirac_status, dirac_site,
                            ganga_status, app_status])

    if pipe_out:
        output(status_list)
    else:
        return status_list

def getStateTime(id, status, pipe_out=True):
    ''' Return the state time from DIRAC corresponding to DIRACJob tranasitions'''
    log = dirac.loggingInfo(id)
    if 'Value' not in log:
        if pipe_out:
            output(None)
            return
        else:
            return None
    L = log['Value']
    checkstr = ''

    if status == 'running':
        checkstr = 'Running'
    elif status == 'completed':
        checkstr = 'Done'
    elif status == 'completing':
        checkstr = 'Completed'
    elif status == 'failed':
        checkstr = 'Failed'
    else:
        checkstr = ''

    if checkstr == '':
        print("%s" % None)
        return

    for l in L:
        if checkstr in l[0]:
            T = datetime.datetime(*(time.strptime(l[3], "%Y-%m-%d %H:%M:%S")[0:6]))
            if pipe_out:
                output(T)
                return
            else:
                return T
    if pipe_out:
        output(None)
    else:
        return None

def getBulkStateTime(job_ids, status, pipe_out=True):
    ''' Function to repeatedly call getStateTime for multiple Dirac Job id and return the result in a dictionary '''
    result = {}
    for this_id in job_ids:
        result[this_id] = getStateTime(this_id, status, pipe_out=False)

    if pipe_out:
        output(result)
    else:
        return result

def monitorJobs(job_ids, status_mapping, pipe_out=True):
    ''' This combines 'status' and 'getBulkStateTime' into 1 function call for monitoring
    '''
    status_info = status(job_ids, status_mapping, pipe_out=False)
    state_job_status = {}
    for job_id, this_stat_info in zip(job_ids, status_info):
        if this_stat_info:
            update_status = this_stat_info[3]
            if update_status not in state_job_status:
                state_job_status[update_status] = []
            state_job_status[update_status].append(job_id)
    state_info = {}
    for this_status, these_jobs in state_job_status.iteritems():
        state_info[this_status] = getBulkStateTime(these_jobs, this_status, pipe_out=False)

    if pipe_out:
        output((status_info, state_info))
    else:
        return (status_info, state_info)

def timedetails(id):
    ''' Function to return the loggingInfo for a DIRAC Job of id'''
    log = dirac.loggingInfo(id)
    d = {}
    for i in range(0, len(log['Value'])):
        d[i] = log['Value'][i]
    output(d)

# DiracAdmin commands
#/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/\/


def getJobPilotOutput(id, dir):
    ''' Get the output of the DIRAC pilot that this job was running on and place it in dir'''
    pwd = os.getcwd()
    try:
        os.chdir(dir)
        os.system('rm -f pilot_%d/std.out && rmdir pilot_%d ' % (id, id))
        result = DiracAdmin().getJobPilotOutput(id)
    finally:
        os.chdir(pwd)
    output(result)


def getServicePorts():
    ''' Get the service ports from the DiracAdmin based upon the Dirac config'''
    output(DiracAdmin().getServicePorts())


def getSitesForSE(se):
    ''' Get the Sites associated with this SE'''
    from DIRAC.Core.Utilities.SiteSEMapping import getSitesForSE
    result = getSitesForSE(storageElement=se)
    output(result)


def getSEsForSite(site):
    ''' Get the list of SE associated with this site'''
    from DIRAC.Core.Utilities.SiteSEMapping import getSEsForSite
    result = getSEsForSite(site)
    output(result)
