#!/bin/env python
#
# check_js_reset - reset check_js state (will cause new job submission)
#
# Author: Marian Babik
# Copyright CERN 2017
#
from __future__ import print_function
import os
import logging
import time
import pickle

import nap.core

log = logging.getLogger('jess')

app = nap.core.Plugin(description='This plugin resets state of check_js.')

app.add_argument('--vo', help='Virtual Organization.', required=True)
app.add_argument('--vo-fqan', help='VOMS primary attribute as FQAN. If given, will be used along with --vo',
                 default='')
app.add_argument('--work-dir', help='Working directory for storing job meta-data, logs, output, etc.',
                 default='/var/lib/gridprobes')
app.add_argument('--backend', help='Job submission backend to be used, options are scondor, condor, cream, arc',
                 required=True)
app.add_argument('--time-limit', help='Time limit on subsequent resets (counted against job file mtime)',
                 default=15, type=int)


@app.metric(metric_name="JobReSubmit")
def test_js_state(args, io):
    fqan_path = None
    if args.vo_fqan and ('/' in args.vo_fqan or '=' in args.vo_fqan):
        fqan_path = args.vo_fqan.replace("/", ".").replace("=", ".")[1:]
    wd = os.path.join(args.work_dir, fqan_path or args.vo, args.backend, args.hostname)
    job_file = os.path.join(wd, 'jobs.pkl')
    args.wd = wd
    if not os.path.exists(wd):
        log.error('Failed to find working directory: {}'.format(wd))
        io.status = nap.CRITICAL
        io.summary = "ETF internal error: failed to find working directory"
        return
    os.chdir(wd)
    log.debug('Working directory: {}'.format(wd))
    if not os.path.isfile(job_file):
        log.debug('Failed to find job file, nothing to reset')
        io.status = nap.OK
        io.summary = "No existing job being tracked, please re-schedule JobState metric"
        return

    with open(job_file, 'r') as jf:
        job_obj = pickle.load(jf)
    if not job_obj.ts_job_start:
        io.status = nap.OK
        io.summary = "No existing job being tracked, please re-schedule JobState metric"
        return

    log.debug('Current start_time difference {}'.format(time.time() - job_obj.ts_job_start))
    if (time.time() - job_obj.ts_job_start) > args.time_limit * 60:
        log.debug('   Removing job file')
        os.remove(job_file)
    else:
        io.status = nap.OK
        io.summary = "New job was submitted recently (less than 15 minutes ago)"
        print("New job was submitted recently, required period between two "
              "submissions is {} minutes to prevent DDoS attacks".format(args.time_limit))
        return
    p_msg = "[{}] SCHEDULE_FORCED_SVC_CHECK;{};{};{}\n".format(time.time(), args.hostname,
                                                               args.prefix + '-JobState-' + args.suffix,
                                                               time.time())
    log.debug('Rescheduling active check: {}'.format(p_msg))
    with open(os.path.abspath('/var/nagios/rw/nagios.cmd'), "w") as cmd_pipe:
        cmd_pipe.write(p_msg)
        cmd_pipe.flush()

    io.status = nap.OK
    io.summary = "New job was submitted, please check JobState metric for progress"


if __name__ == '__main__':
    app.run()
