#!/usr/bin/env python

import os
import sys
import pwd
import logging
import logging.handlers
import time
import re
import shutil
import hashlib
import multiprocessing
import signal

from argparse import ArgumentParser

parser = ArgumentParser(description = 'Schedule routine sequences of Dynamo executables.')
parser.add_argument('sequence', metavar = 'PATH', help = 'Sequence file.')

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

## Load the configuration
from dynamo.dataformat.configuration import Configuration

try:
    config_path = os.environ['DYNAMO_SERVER_CONFIG']
except KeyError:
    config_path = '/etc/dynamo/server_config.json'

config = Configuration(config_path)

## Set up logging (write to stderr unless path is given)
log_level = getattr(logging, config.logging.level.upper())
log_format = '%(asctime)s:%(levelname)s:%(name)s: %(message)s'

LOG = logging.getLogger()
LOG.setLevel(log_level)
if config.logging.get('path', ''):
    log_handler = logging.handlers.RotatingFileHandler(config.logging.path + '/scheduler.log', maxBytes = 1000000, backupCount = 100)
else:
    log_handler = logging.StreamHandler()
LOG.addHandler(log_handler)

# Print some nice logo before we start logging with the timestamp format
LOG.critical('''
+++++++++++++++++++++++++++++++++++++++++++++++
++++++++++++++ DYNAMO SCHEDULER +++++++++++++++
++++++++++++++  v2.0  +++++++++++++++++++++++++
+++++++++++++++++++++++++++++++++++++++++++++++
''')

log_handler.setFormatter(logging.Formatter(fmt = log_format))

## Set the effective user id to config.user
try:
    os.setuid(pwd.getpwnam(config.user).pw_uid)
except OSError:
    LOG.warning('Cannot switch uid to %s (%d).', config.user, pwd.getpwnam(config.user).pw_uid)

## Create the registry
from dynamo.core.registry import DynamoRegistry

registry = DynamoRegistry(config.registry)

result = registry.backend.query('SELECT `id` FROM `users` WHERE `name` = %s', config.scheduler_user)
if len(result) == 0:
    raise RuntimeError('User %s not found' % config.scheduler_user)

user_id = result[0]

## Load the sequences
EXECUTE, WAIT, TERMINATE = range(3)

executables = {} # {title: exec path}
authorized_executables = set() # set of titles
sequences = {} # {name: sequence}
sequence = []

with open(args.sequence) as source:
    iline = -1
    for line in source:
        iline += 1

        line = line.strip()
        if not line or line.startswith('#'):
            continue

        LOG.debug(line)

        # Executable definitions
        # {title} = path  ...  This executable cannot write
        # <title> = path  ...  This executable can be used to write
        matches = re.match('({\S+}|<\S+>)\s*=\s*(\S+)', line)
        if matches:
            enclosed_title = matches.group(1)
            title = enclosed_title[1:-1]
            write_enabled = (enclosed_title[0] == '<')
            executable = matches.group(2)

            # Replace environment variables
            matches = re.findall('\$\(([^\)]+)\)', executable)
            for match in matches:
                executable = executable.replace('$(%s)' % match, os.environ[match])

            try:
                with open(executable) as source:
                    checksum = hashlib.md5(source.read()).hexdigest()
            except:
                LOG.error('Could not open executable %s (line %d).', executable, iline)
                sys.exit(1)

            if write_enabled:
                # Is it really write-enabled?
                sql = 'SELECT `user_id` FROM `authorized_executables` WHERE `title` = %s AND `checksum` = UNHEX(%s)'
                for auth_user_id in registry.backend.query(sql, title, checksum):
                    if auth_user_id == 0 or auth_user_id == user_id:
                        break
                else:
                    LOG.error('Executable %s (%s) is not authorized for server write operation (line %d).', title, executable, iline)
                    sys.exit(1)

            LOG.debug('Define executable %s = %s (write enabled: %d) (line %d)', title, executable, write_enabled, iline)

            executables[title] = executable
            if write_enabled:
                authorized_executables.add(title)
            
            continue

        # Sequence definitions
        # [SEQUENCE title]
        matches = re.match('\[SEQUENCE\s(\S+)\]', line)
        if matches:
            # Sequence header
            LOG.debug('New sequence %s (line %d)', matches.group(1), iline)
            sequence = sequences[matches.group(1)] = []
            continue

        # Sequence executable step definitions
        # {title} options  ...  Read-only execution
        # <title> options  ...  Write-request execution
        matches = re.match('(\*|\+) +({\S+}|<\S+>)\s*(.*)', line)
        if matches:
            critical = (matches.group(1) == '*')
            enclosed_title = matches.group(2)
            title = enclosed_title[1:-1]
            write_request = (enclosed_title[0] == '<')
            arguments = matches.group(3)

            if write_request and title not in authorized_executables:
                LOG.error('Executable %s is not write-enabled (line %d).', iline)

            # Replace environment variables
            matches = re.findall('\$\(([^\)]+)\)', arguments)
            for match in matches:
                arguments = arguments.replace('$(%s)' % match, os.environ[match])

            LOG.debug('Execute %s %s (line %d)', title, arguments, iline)

            sequence.append((EXECUTE, title, arguments, critical, write_request))
            continue

        matches = re.match('WAIT\s+(.*)', line)
        if matches:
            try:
                sleep_time = eval(matches.group(1))
                if type(sleep_time) not in [int, float]:
                    raise RuntimeError()
            except:
                LOG.error('Wait time %s is not a numerical expression (line %d).', matches.group(1), iline)
                sys.exit(1)

            LOG.debug('Wait for %d seconds (line %d)', sleep_time, iline)
            sequence.append((WAIT, sleep_time))
            continue

        if line == 'TERMINATE':
            sequence.append([TERMINATE])

if len(sequences) == 0:
    LOG.error('No sequence was loaded.')
    sys.exit(1)

## Create the executable directory and copy the scripts over
for name, sequence in sequences.iteritems():
    work_dir = config.scheduler_path + '/' + name
    try:
        shutil.rmtree(work_dir)
    except OSError:
        pass
    os.makedirs(work_dir)
    os.chmod(work_dir, 0777)

    with open(work_dir + '/log.out', 'w'):
        pass

    os.chmod(work_dir + '/log.out', 0666)

    with open(work_dir + '/log.err', 'w'):
        pass

    os.chmod(work_dir + '/log.err', 0666)

    for action in sequence:
        if action[0] == EXECUTE:
            title = action[1]

            path = '%s/%s' % (work_dir, title)
            if os.path.exists(path):
                # This executable is used multiple times in the sequence
                continue

            executable = executables[title]

            os.makedirs(path)
            shutil.copy(executable, path + '/exec.py')

            # Make symlinks from subdir/_std(out|err) to log.(out|err)
            os.symlink(work_dir + '/log.out', path + '/_stdout')
            os.symlink(work_dir + '/log.err', path + '/_stderr')

LOG.info('Created working directories under %s.', config.scheduler_path)

def loop_sequence(sequence_name, sequence, user_id, config):
    ## Ignore SIGINT - see note above proc.terminate() at the end of the file
    ## We will react to SIGTERM by raising KeyboardInterrupt
    from dynamo.utils.signaling import SignalConverter

    signal.signal(signal.SIGINT, signal.SIG_IGN)

    ## Use the same logger as the main process, but with a different format
    LOG.info('Sequence %s (PID: %d)', sequence_name, os.getpid())

    log_format = '%(asctime)s:%(levelname)s:{sequence} ({pid}): %(message)s'.format(sequence = sequence_name, pid = os.getpid())
    LOG.handlers[0].setFormatter(logging.Formatter(fmt = log_format))

    ## Create a registry handle - we'll be submitting execution requests from the backend
    registry = DynamoRegistry(config.registry)

    ## SQL templates
    insert_sql = 'INSERT INTO `action` (`user_id`, `write_request`, `title`, `path`, `args`)'
    insert_sql += ' VALUES ({user_id}, %s, %s, %s, %s)'.format(user_id = user_id)

    poll_sql = 'SELECT `status` FROM `action` WHERE `id` = %s'

    work_dir = config.scheduler_path + '/' + sequence_name
    title = ''
    task_id = 0

    try:
        ## Start an infinite loop
        while True:
            for action in sequence:
                if action[0] == EXECUTE:
                    title, arguments, critical, write_request = action[1:]
                    executable = executables[title]
    
                    LOG.info('Starting %s', title)
                    LOG.info('Command: %s %s', executable, arguments)
    
                    with open(work_dir + '/log.out', 'a') as out:
                        out.write('------------------------ ' + title + ' ------------------------\n')
                        out.write('%s %s\n\n' % (executable, arguments))
    
                    with open(work_dir + '/log.err', 'a') as out:
                        out.write('------------------------ ' + title + ' ------------------------\n\n')
    
                    # Submit the task. Actual executable is in work_dir/title
                    path = '%s/%s' % (work_dir, title)
                    task_id = registry.backend.query(insert_sql, write_request, title, path, arguments)
    
                    # Poll for task completion
                    while True:
                        time.sleep(1)
                        result = registry.backend.query(poll_sql, task_id)
                        if len(result) == 0:
                            # len(result) == 0 -> entry disappeared!?
                            status = 'unknown'
                            break
                        else:
                            status = result[0]
                            if status not in ['new', 'run']:
                                break

                    task_id = 0
    
                    try:
                        with open(work_dir + '/log.out', 'a') as out:
                            out.write('\n')
                    except:
                        pass
    
                    try:
                        with open(work_dir + '/log.err', 'a') as out:
                            out.write('\n')
                    except:
                        pass
    
                    LOG.info('Finished %s with (status: %s)\n', title, status)

                    if critical and status != 'done':
                        LOG.error('Critical executable %s failed. Terminating sequence %s.', title, sequence_name)
                        ## We don't have a way to restart individual sequences - once this happens, the entire scheduler has to be restarted..
                        break
    
                elif action[0] == WAIT:
                    LOG.info('Sleeping for %d seconds.\n', action[1])
                    time.sleep(action[1])
    
                elif action[0] == TERMINATE:
                    LOG.info('Terminating sequence.')
                    break
    
            else:
                # Passed through the sequence without breaking
                continue
    
            # There was a break -> sequence terminated
            break

    except KeyboardInterrupt:
        if task_id != 0:
            # Task aborted - update the status from the backend so Dynamo picks it up
            registry.backend.query('UPDATE `action` SET `status` = \'killed\' WHERE `id` = %s', task_id)
            LOG.info('Process interrupted. aborting task %s.', title)
        else:
            LOG.info('Process interrupted.')


## Prepare to catch a kill
from dynamo.utils.signaling import SignalConverter
signal_converter = SignalConverter(logger = LOG)
signal_converter.set(signal.SIGTERM)

try:
    ## Create a process for each sequence and start as daemons
    processes = []
    for name, sequence in sequences.iteritems():
        proc = multiprocessing.Process(target = loop_sequence, name = name, args = (name, sequence, user_id, config))
        proc.daemon = True
        proc.start()
        LOG.info('Started sequence %s.', name)
        
        processes.append(proc)

    ## Sleep indefinitely
    for proc in processes:
        proc.join()

except KeyboardInterrupt:
    LOG.info('Scheduler terminated.')

except:
    LOG.error('Exception in scheduler. Terminating all sequences.')
    raise

finally:
    # This can be a Ctrl+C or a kill.
    # Ctrl+C will pass SIGINT to all child processes (if this process is the head of the
    # foreground process group). In this case calling terminate() will duplicate signals
    # in the child. Child processes have to always ignore SIGINT and be killed only from
    # SIGTERM sent by the line below.
    for proc in processes:
        proc.terminate()
        proc.join(2)

LOG.info('All sequences completed.')
