#!/usr/bin/env python

###############################################################
####### This script will spit out png files monitoring ########
####### the copy status through Phedex on three levels: #######
####### -per replica, -per request, -per site #################
#
####### yiiyama@mit.edu, bmaier@mit.edu #######################
###############################################################

import sys
import os
import time
import shutil
import rrdtool
import selinux
import collections
import csv

from glob import glob
from datetime import datetime, timedelta

from argparse import ArgumentParser

parser = ArgumentParser(description = 'Track transfers')
parser.add_argument('--config', '-c', metavar = 'CONFIG', dest = 'config', required = True, help = 'Configuration JSON.')

args = parser.parse_args()
sys.argv = []

from dynamo.dataformat import Configuration
import dynamo.history.impl as history_impl
import dynamo.operation.impl as operation_impl
from dynamo.utils.parallel import Map
from dynamo.core.executable import make_standard_logger

config = Configuration(args.config)

LOG = make_standard_logger(config.log_level)

rrd_dir = config.rrd_path_base + '/track_transfers'
try:
    os.makedirs(rrd_dir)
except:
    pass

sites = glob(rrd_dir+'/*')
try:
    sites.remove(rrd_dir + '/total.rrd')
    sites.remove(rrd_dir + '/monitoring')
except:
    pass

history = getattr(history_impl, config.history.module)(config.history.config)

sites_with_open_transfers = []
records = collections.defaultdict(set)

sites_ongoings = []
sites_total = [] # only requests with copied != total
sites_copied = [] # only requests with copied != total
sites_total_stuck = [] # only datasets with copied != total
sites_copied_stuck = [] # only datasets with copied != total
sites_total_ongoing = [] # only datasets with copied != total
sites_copied_ongoing = [] # only datasets with copied != total

# Get all sites with ongoing transfers

for partition in config.partitions:
    partition_records = history.get_incomplete_copies(partition)    
    for record in partition_records:
        site = history.get_site_name(record.operation_id)
        records[site].add(record)
        if site not in sites_with_open_transfers:
            sites_with_open_transfers.append(site)

#LOG.debug('Incomplete copies at the moment: %d', len(request_ids))

copy = getattr(operation_impl, config.copy_op.module)(config.copy_op.config)

interval = int(config.rrd_interval)

timestamp = int(time.time()) / interval * interval

incomplete_replicas_rrd = []

# Keeping track of total volume ...
total_volume = 0
copied_volume = 0
total_rrdfile = rrd_dir+'/total.rrd'


start = (int(time.time()) / interval - 1) * interval

if not os.path.exists(total_rrdfile):
    rrdtool.create(total_rrdfile, '--start', str(start), '--step', str(interval),
                   'DS:copied:GAUGE:%d:0:U' % (interval * 800),
                   'DS:total:GAUGE:%d:0:U' % (interval * 800),
                   'RRA:LAST:0:1:%i' % 1344    )

# Now create new rrd files and create pngs
def is_transfer_stuck(rrd_file):
    stuck = 0
    result = rrdtool.fetch(rrd_file, "LAST")

    lasttime = result[0][1]
    firsttime = lasttime - 6*24*3600

    result = rrdtool.fetch(rrd_file, "LAST", '-s', str(firsttime), '-e', str(lasttime))

    rows = result[2]
    try:
        while rows[len(rows)-1][1] is None:
            rows = rows[:-1]

        if len(rows) > 480 and (rows[len(rows)-1][0] - rows[len(rows)-481][0])/rows[len(rows)-1][1] < 0.01:
        #480 corresponds to 5 days * 24 hours * 4 timestamps per hour
            stuck = 1
    except:
        pass

    return stuck

def exec_get(sitename):
    global rrd_dir
    global records
    global copy
    global incomplete_replicas_rrd
    global interval

    global sites_ongoings
    global sites_total
    global sites_copied
    global sites_total_stuck
    global sites_copied_stuck
    global sites_total_ongoing
    global sites_copied_ongoing

    # Need this here for some reason
    import rrdtool
    import csv

    debug_counter = 0

    ongoings = []
    site_total = 0 # request level
    site_copied = 0 # request level
    site_total_stuck = 0 # dataset level
    site_copied_stuck = 0 # dataset level
    site_total_ongoing = 0 # dataset level
    site_copied_ongoing = 0 # dataset level

    rrd_filepath = rrd_dir + '/' + sitename

    if not os.path.exists(rrd_filepath):
        # Create path corresponding to site
        os.mkdir(rrd_filepath)

    try:
        os.unlink("%s/filelist.txt" % rrd_filepath)
    except:
        pass

    for record in records[sitename]:
        debug_counter += 1
        #if debug_counter > 2:
        #    break    

        request_id = record.operation_id

        LOG.debug(request_id)

        status = copy.copy_status(request_id)

        request_total = 0
        request_copied = 0

        if len(status) == 0:
            record.completed = 1 #old transfer. seems to not exist anymore.

        for (site, dataset), (total, copied, last_update) in status.items():
            LOG.debug(dataset)

            if copied is not None:# That happens sometimes for very, very, very old requests
                site_total += total
                site_copied += copied

            tmp = rrd_filepath + '/'  +  str(request_id) + '_' + dataset.replace('/', '+') + '.rrd'
            rrd_file = tmp.replace('+','',1)
            incomplete_replicas_rrd.append(rrd_file)
         
            if not os.path.exists(rrd_file) and copied != total:
                # RRD does not exist yet
                start = (int(time.time()) / interval - 1) * interval

                # Write rrd file
                rrdtool.create(rrd_file, '--start', str(start), '--step', str(interval),
                               'DS:copied:GAUGE:%d:0:U' % (interval * 800),
                               'DS:total:GAUGE:%d:0:U' % (interval * 800),
                               'RRA:LAST:0:1:%i' % 1344    )

                # data source
                #  DS:<name>:<type>:<heartbeat>:<min>:<max>
                #  type = GAUGE: quantity that has a value at each time point
                #  heartbeat: "maximum number of seconds that may pass between two updates of this data source before the value of the data source is assumed to be *UNKNOWN*"
                #  min/max = U: unknown
                # round robin archive (RRA)
                #  RRA:<type>:<xff>:<nsteps>:<nrows>
                #  type = LAST: just use the last value, no averaging etc.
                #  xff: fraction of <nsteps> that can have UNKNOWN as the value
                #  nsteps: number of steps used for calculation
                #  nrows: number of records to keep

                # change selinux context of the RRD so that it can be read by a apache-invoked PHP script
                try:
                    selinux.chcon(rrd_file, 'unconfined_u:object_r:httpd_sys_content_t:s0')
                except:
                    pass

            try:
                # Keeping track of the request status
                request_total += total
                request_copied += copied            
                if copied != total:
                    is_stuck = is_transfer_stuck(rrd_file)
                    ongoings.append([request_id,dataset,copied,total,is_stuck])
                    site_total_stuck += is_stuck*total
                    site_copied_stuck += is_stuck*copied
                    site_total_ongoing += total
                    site_copied_ongoing += copied
                    rrdtool.update(rrd_file, '%d:%d:%d' % (timestamp, copied, total))
                else:
                    try:
                        os.unlink(rrd_file)
                    except:
                        pass
            except:
                pass
            
        # Update history DB
        if len(status) != 0:
            if request_total != record.size or request_copied == request_total:
                record.size = request_total
                record.completed = (request_copied == request_total)        
                #history.update_copy_entry(record)


    sites_ongoings.append([sitename,len(ongoings)])
    sites_total.append([sitename,site_total])
    sites_copied.append([sitename,site_copied])
    sites_total_stuck.append([sitename,site_total_stuck])
    sites_copied_stuck.append([sitename,site_copied_stuck])
    sites_total_ongoing.append([sitename,site_total_ongoing])
    sites_copied_ongoing.append([sitename,site_copied_ongoing])

    with open("%s/filelist.txt" % rrd_filepath, "w") as csvfilelist:
        fieldnames = ["id","name","copied","total","stuck"]
        writer = csv.DictWriter(csvfilelist, fieldnames=fieldnames)
        writer.writerow(dict(zip(fieldnames,fieldnames)))
        for ongoing in ongoings:
            writer.writerow({"id" : ongoing[0], "name" : ongoing[1], "copied" : ongoing[2], "total" : ongoing[3], "stuck" : ongoing[4]})        


Map(config.parallel).execute(exec_get, sites_with_open_transfers)

for site in sites_with_open_transfers:
    for record in records[site]:
        history.update_copy_entry(record)



# Creating overview files
try:
    os.unlink("%s/overview.txt" % rrd_filepath)
except:
    pass

with open("%s/overview.txt" % rrd_dir, "w") as overview:
    fieldnames = ["sitename","ongoing","total","copied","total_stuck","copied_stuck"]
    writer = csv.DictWriter(overview, fieldnames=fieldnames)
    writer.writerow(dict(zip(fieldnames,fieldnames)))
    for site in sites_with_open_transfers:
        site_total = 0
        for x in sites_total:
            if x[0] == site:
                site_total += x[1]

        if site_total == 0:
            continue

        site_ongoing = 0
        for x in sites_ongoings:
            if x[0] == site:
                site_ongoing += x[1]

        site_total_ongoing = 0
        for x in sites_total_ongoing:
            if x[0] == site:
                site_total_ongoing += x[1]

        site_copied_ongoing = 0
        for x in sites_copied_ongoing:
            if x[0] == site:
                site_copied_ongoing += x[1]

        site_total_stuck = 0
        for x in sites_total_stuck:
            if x[0] == site:
                site_total_stuck += x[1]

        site_copied_stuck = 0
        for x in sites_copied_stuck:
            if x[0] == site:
                site_copied_stuck += x[1]

        writer.writerow({"sitename" : site, "ongoing" : site_ongoing, "total" : site_total_ongoing,
                         "copied" : site_copied_ongoing, "total_stuck" : site_total_stuck, "copied_stuck" : site_copied_stuck})

total_volume = sum((x[1] for x in sites_total))
copied_volume = sum((x[1] for x in sites_copied))

try:
    rrdtool.update(total_rrdfile, '%d:%d:%d' % (timestamp, copied_volume, total_volume))
except:
    pass


# Deletion part - first delete rrd files of completed requests that are older than one week, since we do not want them to be a part of the graphs anymore 
for site in sites:
    try:
        existing_rrds = glob(site+'/*.rrd')
        
        older_than = datetime.now() - timedelta(days=20)
        
        for existing_rrd in existing_rrds:
            filetime = datetime.fromtimestamp(os.path.getmtime(existing_rrd))
            if existing_rrd not in incomplete_replicas_rrd and filetime < older_than:
                # Delete pngs and rrd files
                os.unlink(existing_rrd)

    except:
        pass

# Copying rrds to the /var/www location
target_dir = config.rrd_publish_target + '/monitoring'
try:
    os.makedirs(target_dir)
except OSError:
    pass

for entry in os.listdir(target_dir):
    if entry.startswith('T'):
        target = target_dir + '/' + entry
        try:
            shutil.rmtree(target)
        except:
            pass

try:
    os.unlink(target_dir + '/total.rrd')
except:
    pass

try:
    os.unlink(target_dir + '/overview.txt')
except:
    pass

for entry in os.listdir(rrd_dir):
    if entry.startswith('T'):
        source = rrd_dir + '/' + entry
        target = target_dir + '/' + entry
        shutil.copytree(source, target)

shutil.copy(rrd_dir + '/total.rrd', target_dir)
shutil.copy(rrd_dir + '/overview.txt', target_dir)
