#!/usr/bin/python

import os
import sys
import re
import requests
import json
import subprocess

def get_cpu_count():
    # cpuset way is a right way
    try:
        cpuset_cpus = re.search(r'(?m)^Cpus_allowed:\s*(.*)$', open('/proc/self/status').read())
        if cpuset_cpus:
            cpus = bin(int(cpuset_cpus.group(1).replace(',', ''), 16)).count('1')
            if cpus > 0:
                return cpus
    except IOError:
        pass
    # if old Linux just count /proc/cpuinfo
    try:
        cpus = open('/proc/cpuinfo').read().count('processor\t:')
        if cpus > 0:
            return cpus
    except IOError:
        pass
    # return 1 CPU in case of previous methods failed
    return 1

def get_memory_ammount():
    # try MemAvailable
    mem_avail = re.search(r'(?m)^MemAvailable:\s*(.*)\skB$', open('/proc/meminfo').read())
    if mem_avail:
        return int(mem_avail.group(1))/1024

    # use MemTotal in case of old kernel
    mem_avail = re.search(r'(?m)^MemTotal:\s*(.*)\skB$', open('/proc/meminfo').read())
    return int(mem_avail.group(1))/1024

def get_disk_space(path):
    if os.path.isdir(path):
        s = os.statvfs(path)
    else:
        s = os.statvfs('/')
    return s.f_bsize * s.f_bavail / 1024 / 1024


def detect_gpu ():
    try:
        output = subprocess.check_output('nvidia-smi -L', shell=True)
    except subprocess.CalledProcessError:
        return 0
    else:
        return output.count('\n')

def detect_rtes (rte_dir):
    if os.path.isdir(rte_dir):
        return os.listdir(rte_dir)
    else:
        return None


def send_request (data):
    try:  
        gloria_url = os.environ["GLORIA_URL"] + '/pilot.php'
        gloria_key = os.environ["GLORIA_KEY"]
    except KeyError: 
        print 'ERROR: Gloria access parameters should be defined in the environment'
        sys.exit(1)

    data['auth'] = {'type': 'key', 'key': gloria_key}

    requests.packages.urllib3.disable_warnings()
    r = requests.put(gloria_url, data=json.dumps(data), verify=False)
    if r.status_code == 200:
        return r.json()
    else:
        return False
   
def stage_data_in (url, name):
    cmd =  'arccp ' + url + ' ' + name
    return subprocess.call(cmd, shell=True)

def stage_data_out (url, name):
    cmd =  'arccp ' + name + ' ' + url
    return subprocess.call(cmd, shell=True)

#
# Get configuration from environment
#
workdir = os.getenv('GLORIA_PILOT_WORKDIR', '/mnt/data')
rte_dir  = os.getenv('GLORIA_RTE_DIR', '/etc/gloria/RTE')

# Go to working directory
if not os.path.isdir(workdir):
    os.mkdir(workdir)
os.chdir(workdir)

# 
# Prepare request
# 
data = {}
data['request_type'] = 'fetch-job'
data['resources'] = { 'memory': get_memory_ammount(),
                      'cpus': get_cpu_count(),
                      'disk': get_disk_space(workdir) }

data['runtime_features'] = []
gpus = detect_gpu()
if gpus:
    data['runtime_features'].append('GPU')

rtes = detect_rtes(rte_dir)
if rtes:
    data['runtime_software'] = [ r for r in rtes ]

#
# Send request to Gloria
#
job = send_request(data)
if not job:
    print "FATAL: Failed to fetch job from Gloria"
    sys.exit(1)

# Save proxy to file
if 'proxy' in job.keys():
    proxy_f = open('x509_job_proxy','w')
    proxy_f.write(job['proxy'])
    proxy_f.close()
    os.chmod('x509_job_proxy', 384)
    os.environ['X509_USER_PROXY'] = workdir + '/x509_job_proxy'

# Process stage-in
for inf in job['input_files']:
    for in_file,url in inf.iteritems():
        stage_data_in(url,in_file)

# Set environment
for env in job['environment']:
    for name,value in env.iteritems():
            os.environ[name] = value

# Source runtime
job_cmd = ''
for rte in job['runtime_software']:
    rtefile = rte_dir + '/' + rte
    if os.path.isfile(rtefile):
        job_cmd = job_cmd + 'source ' + rtefile + '; '

# Execute job script
job_script = job['jobinfo']['executable']
job_cmd = job_cmd + '/bin/bash ' + job_script
subprocess.call(job_cmd, shell=True)

# Process stage-out
for of in job['output_files']:
    for o_file,url in of.iteritems():
        stage_data_out(url,o_file)