#!/usr/bin/env python

# Copyright (C) 2012-2013 Science and Technology Facilities Council.
# Copyright (C) 2016-2021 East Asian Observatory.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from codecs import latin_1_encode, latin_1_decode
import os
import sys
import time
from optparse import OptionParser

from crab import CrabError, CrabStatus
from crab.client import CrabClient
from crab.util.compat import subprocess, subprocess_options, \
    subprocess_call, subprocess_communicate, TimeoutExpired
from crab.util.pid import pidfile_write, pidfile_running, pidfile_delete
from crab.util.string import split_crab_vars, true_string


def main():
    parser = OptionParser()
    parser.add_option(
        '-c',
        type='string', dest='command',
        help='specify the COMMAND to execute', metavar='COMMAND')
    parser.add_option(
        '--id',
        type='string', dest='crabid',
        help='set Crab job ID', metavar='ID')
    parser.add_option(
        '--shell',
        type='string', dest='shell',
        help='use SHELL to execute COMMAND', metavar='SHELL')
    parser.add_option(
        '--pidfile',
        type='string', dest='pidfile',
        help='use PIDFILE to avoid re-running COMMAND', metavar='PIDFILE')

    (options, args) = parser.parse_args()

    # Determine command to execute

    if len(args) != 0:
        parser.error('no arguments required')
    if options.command is None:
        parser.error('COMMAND not specified')

    (command, vars) = split_crab_vars(options.command)

    # Update environment with parsed variables and extract any
    # additional variables from the environment.

    env = os.environ

    for envvar in vars:
        env[envvar] = vars[envvar]
    for envvar in env:
        if envvar.startswith('CRAB') and envvar not in vars:
            vars[envvar] = env[envvar]

    # Determine shell to use
    #
    # Note that cron defaults to sh regardless of the user's shell.

    shell = '/bin/sh'

    if options.shell:
        shell = options.shell
    elif 'CRABSHELL' in vars:
        shell = vars['CRABSHELL']

    # Check CRABIGNORE variable from command if present.

    ignore = ('CRABIGNORE' in vars and true_string(vars['CRABIGNORE']))

    # Look for cron job ID
    #
    # We could leave the CRABID in the command for the shell to
    # extract, but stripping it off makes sure the command matches
    # that in the database and allows for shells which cannot handle
    # variables given at the start of a command.

    crabid = None

    if options.crabid:
        crabid = options.crabid
    elif 'CRABID' in vars:
        crabid = vars['CRABID']

    # Prepare Crab client options -- will be passed to client constructor
    # when needed.

    client_options = {'command': command, 'crabid': crabid}

    # Check for presence of PID file, if a PID file has been configured,

    pidfile = None

    if options.pidfile:
        pidfile = options.pidfile
    elif 'CRABPIDFILE' in vars:
        pidfile = vars['CRABPIDFILE']

    if pidfile is not None:
        if pidfile_running(pidfile):
            # Only report "already running" status when not in "ignore" mode.
            if not ignore:
                try:
                    client = CrabClient(**client_options)
                    client.finish(CrabStatus.ALREADYRUNNING)
                except CrabError:
                    err = sys.exc_info()[1]
                    print('crabsh: ' + command)
                    print('Failed to notify job already running.')
                    print('ERROR: ' + str(err) + '\n')

            return 0

        # Write PID file based on our PID.  (Changed in version 0.5.0:
        # previously we wrote the PID of the child process.)
        pidfile_write(pidfile, os.getpid())

    # Check for watchdog timeout variable.

    watchdog_timeout = None

    if 'CRABWATCHDOG' in vars:
        watchdog_timeout = 60 * int(vars['CRABWATCHDOG'])

    # Begin try-finally block to try to ensure PID file is deleted at end.

    try:

        # In "ignore" mode, simply run the command directly (i.e. without
        # redirecting stdout and stderr).

        if ignore:
            try:
                return subprocess_call([shell, '-c', command], env=env,
                                       timeout=watchdog_timeout,
                                       **subprocess_options)

            except OSError:
                err = sys.exc_info()[1]
                print('crabsh (' + shell + '): ' + command)
                print('ERROR: ' + str(err))

                return 1

            except TimeoutExpired:
                print('crabsh (' + shell + '): ' + command)
                print('Killed by watchdog.')

                return 1

        # Otherwise attempt to execute the command with notifications sent
        # to the Crab server.

        client = CrabClient(**client_options)

        try:
            response = client.start()

            # If the server sent an inhibit response, check the config to
            # see whether crabsh.allow_inhibit is on or not.
            if response.get('inhibit', False):
                allow_inhibit = True
                try:
                    allow_inhibit = true_string(
                            client.config.get('crabsh', 'allow_inhibit'))
                except:
                    # The configuration key might have been missing -- since
                    # this is a crabsh-specific option, the crab client will
                    # not have ensured it is present.
                    pass

                if allow_inhibit:
                    try:
                        client.finish(CrabStatus.INHIBITED)
                    # except CrabError as err:
                    except CrabError:
                        err = sys.exc_info()[1]
                        print('crabsh: ' + command)
                        print('Failed to acknowledge inhibited job.')
                        print('ERROR: ' + str(err) + '\n')

                    return 0

        # except CrabError as err:
        except CrabError:
            err = sys.exc_info()[1]
            print('crabsh: ' + command)
            print('Failed to notify job start.')
            print('ERROR: ' + str(err) + '\n')

        returncode = None

        try:
            p = subprocess.Popen([shell, '-c', command],
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 env=env,
                                 **subprocess_options)

            try:
                (stdoutdata, stderrdata) = subprocess_communicate(
                    p, timeout=watchdog_timeout)
                returncode = p.returncode

            except TimeoutExpired:
                p.kill()
                (stdoutdata, stderrdata) = subprocess_communicate(p)

            stdoutdata = (latin_1_decode(stdoutdata, 'replace'))[0]
            stderrdata = (latin_1_decode(stderrdata, 'replace'))[0]

            status = CrabStatus.SUCCESS

            if returncode is None:
                status = CrabStatus.WATCHDOG

            elif returncode:
                status = CrabStatus.FAIL

            client.finish(status, stdoutdata, stderrdata)

        # except OSError as err:
        except OSError:
            err = sys.exc_info()[1]
            try:
                client.finish(CrabStatus.COULDNOTSTART, str(err))
            except CrabError:
                print('crabsh (' + shell + '): ' + command)
                print('Failed to notify that job could not start.')
                print('ERROR: ' + str(err))

        # except CrabError as err:
        except CrabError:
            err = sys.exc_info()[1]
            # Print fall-back message for cron to send by email (to the
            # crontab owner or address set its MAILTO variable.
            print('crabsh: ' + command)
            print('Failed to notify job finish.')
            print('ERROR: ' + str(err))
            if returncode is None:
                print('\nKilled by watchdog timeout.')
            else:
                print('\nRETURN CODE: ' + str(returncode))
            if stdoutdata != '':
                print('\nSTDOUT:')
                print(stdoutdata)
            if stderrdata != '':
                print('\nSTDERR:')
                print(stderrdata)

        else:
            # Echo the output only if we didn't already print it due
            # to an exception occurring.
            if ('CRABECHO' in vars) and true_string(vars['CRABECHO']):
                if stdoutdata != '':
                    print(stdoutdata)
                if stderrdata != '':
                    if stdoutdata != '':
                        print('\n\nStandard Error:\n')
                    print(stderrdata)

    finally:
        if pidfile is not None:
            pidfile_delete(pidfile)


if __name__ == "__main__":
    main()
