#!/usr/bin/env python

# Copyright (C) 2012 Science and Technology Facilities Council.
#
# 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 subprocess
import time
from optparse import OptionParser

from crab import CrabError, CrabStatus
from crab.client import CrabClient
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.

    if 'CRABIGNORE' in vars and true_string(vars['CRABIGNORE']):
        try:
            return subprocess.call([shell, '-c', command], env=env)

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

            return 1

    # 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']

    # Look for PID file setting

    pidfile = None

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

    # Attempt to execute the command

    client = CrabClient(command, crabid=crabid)

    if pidfile is not None:
        if pidfile_running(pidfile):
            try:
                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

    try:
        client.start()
    #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')

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

        if pidfile is not None:
            # Add a short delay before creating the PID file in case
            # the cmdline which pidfile_write uses changes when the
            # shell starts running the command.
            time.sleep(5)
            pidfile_write(pidfile, p.pid)

        (stdoutdata, stderrdata) = p.communicate()

        if pidfile is not None:
            pidfile_delete(pidfile)

        returncode = CrabStatus.SUCCESS

        if p.returncode:
            returncode = CrabStatus.FAIL

        client.finish(returncode,
                      (latin_1_decode(stdoutdata, 'replace'))[0],
                      (latin_1_decode(stderrdata, 'replace'))[0])

    #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))
        print('\nRETURN CODE: ' + str(returncode))
        print('\nSTDOUT:')
        print(stdoutdata)
        print('\n\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)

if __name__ == "__main__":
    main()
