#!/usr/bin/env python
#
# run-as-hook - runs a script or command as if it were a hook, on this computer
"""
Examples:

$ PASSWORD=$(jitsu run-as-hook mysql/0 relation-get -r db:0 password)

Note that a relation id must be supplied for relation hook commands.

Define a script refresh.sh like so::

  #!/bin/bash
  for relation_id in $(relation-ids db); do
      relation-set -r $relation_id ready=$1
  done

$ jitsu run-as-hook --loglevel=DEBUG mysql/0 refresh.sh 32
2012-04-24 11:02:19,671 juju.ec2:WARNING txaws.client.ssl unavailable for SSL hostname verification
2012-04-24 11:02:19,671 juju.ec2:WARNING EC2 API calls encrypted but not authenticated
2012-04-24 11:02:19,671 juju.ec2:WARNING S3 API calls encrypted but not authenticated
2012-04-24 11:02:19,671 juju.ec2:WARNING Ubuntu Cloud Image lookups encrypted but not authenticated
2012-04-24 11:02:19,673 juju.common:INFO Connecting to environment...
2012-04-24 11:02:20,547 juju.common:DEBUG Connecting to environment using ec2-50-18-146-37.us-west-1.compute.amazonaws.com...
2012-04-24 11:02:20,548 juju.state.sshforward:DEBUG Spawning SSH process with remote_user="ubuntu" remote_host="ec2-50-18-146-37.us-west-1.compute.amazonaws.com" remote_port="2181" local_port="60629".
2012-04-24 11:02:22,186 juju.common:DEBUG Environment is initialized.
2012-04-24 11:02:22,186 juju.common:INFO Connected to environment.
2012-04-24 11:02:22,592 jitsu-run-as-hook-output:DEBUG Cached relation hook contexts: ['db:0', 'db:1']
2012-04-24 11:02:23,352 juju.jitsu.do.unit-agent:DEBUG Setting relation db:0
2012-04-24 11:02:23,794 juju.jitsu.do.unit-agent:DEBUG Setting relation db:1
2012-04-24 11:02:23,875 jitsu-run-as-hook-output:DEBUG hook jitsu-run-as-hook-script exited, exit code 0.
2012-04-24 11:02:24,666 jitsu-run-as-hook-output:DEBUG Flushed values for hook 'jitsu-run-as-hook-script'
    Setting changed: 'ready'=u'32' (was '42') on 'db:0'
    Setting changed: 'ready'=u'32' (was '42') on 'db:1'
"""

from StringIO import StringIO
import argparse
import zookeeper
import logging
import os
import shutil
import sys
import tempfile
import traceback

from twisted.internet import reactor
from twisted.internet.defer import inlineCallbacks

from aiki.invoker import InvokerUnitAgent
from aiki.introspect import get_zk_client_connector
import juju


def main():
    loglevels=dict(
            CRITICAL=logging.CRITICAL,
            ERROR=logging.ERROR,
            WARNING=logging.WARNING,
            INFO=logging.INFO,
            DEBUG=logging.DEBUG)

    parser = argparse.ArgumentParser(
        description="runs a script locally with the context of the given unit")
    parser.add_argument(
        "-e", "--environment", default=None,
        help="Environment to act upon (otherwise uses default)")
    parser.add_argument(
        "--loglevel", default=None, choices=loglevels,
        help="Log level",
        metavar="CRITICAL|ERROR|WARNING|INFO|DEBUG")
    parser.add_argument(
        "--verbose", "-v", default=False, action="store_true",
        help="Enable verbose logging")
    parser.add_argument("unit_name", help="Local unit name for running command", metavar="SERVICE_UNIT")
    parser.add_argument("command", help="Command", metavar="COMMAND")
    parser.add_argument("arg", nargs="*", help="Optional args", metavar="ARG")

    # Collect additional args to be passed through, such as arguments like -r db:0
    options, extra = parser.parse_known_args()
    extra.extend(options.arg)
    options.extra = extra

    # Prefix command with path if necessary
    if os.path.exists(os.path.join(os.getcwd(), options.command)):
        options.command = os.path.abspath(options.command)

    # Set logging defaults with respect to verbose
    options.loglevel = loglevels.get(options.loglevel)
    if options.verbose:
        if options.loglevel is None:
            options.loglevel = logging.DEBUG
    else:
        # Avoid potentially voluminous ZK debug logging 
        zookeeper.set_debug_level(0)
        if options.loglevel is None:
            options.loglevel = logging.WARNING

    # Some extra complexity here in logging level setup vs a normal
    # root logger; see the comments on the jitsu-run-as-hook-output logging
    # for the motivation on why.
    log_options = {
        "level": min(options.loglevel, logging.INFO),
        "format": "%(asctime)s %(name)s:%(levelname)s %(message)s"
        }
    logging.basicConfig(**log_options)
    juju_root_log = logging.getLogger("juju")  # NB Cannot do this with the root logger!
    juju_root_log.setLevel(logging.getLevelName(options.loglevel))

    # Introspect to find the right client connector; this will be
    # called, then yielded upon, to get a ZK client once async code is
    # entered
    try:
        connector = get_zk_client_connector(options)
    except Exception, e:
        print >> sys.stderr, e
        sys.exit(1)

    # Setup temporary unit agent to invoke command + any args as a
    # hook and run it
    temp_dir = tempfile.mkdtemp(prefix="jitsu-run-as-hook-")
    result = [0]
    reactor.callWhenRunning(
        run_one, juju_do, options.verbose, result,
        connector, temp_dir, options)
    reactor.run()
    shutil.rmtree(temp_dir)
    sys.exit(result[0])


@inlineCallbacks
def run_one(func, verbose, result, *args, **kw):
    try:
        yield func(result, *args, **kw)
    except Exception, e:
        result[0] = 1
        if verbose:
            traceback.print_exc()  # Writes to stderr
        else:
            print >> sys.stderr, e
    finally:
        reactor.stop()


@inlineCallbacks
def juju_do(result, connector, temp_dir, options):
    """Runs the hook in the reactor, as if it were a unit agent.

    `result` is a list of one object that is used to communicate back
    the exit code. Otherwise stdout/stderr are used for the unit agent
    output and error text, if any.
    """

    # Most of this code is to handle the many, many ways that the hook
    # can fail, while capturing useful exit codes and and error text
    # (and omitting useless junk).

    # Connect to ZK, either from admin machine or on a Juju machine
    client = yield connector()

    # Unit agent output/error
    log_file = StringIO()
    error_file = StringIO()

    # And actually run the hook
    invoker_ua = InvokerUnitAgent(client, temp_dir, log_file, error_file, options)
    yield invoker_ua.start()
    try:
        result[0] = yield invoker_ua.run_hook(options.command, options.extra)
    except juju.errors.CharmInvocationError, e:
        # Get exit code from this exception, but pick up the actual
        # error text from the error_file, since it is more detailed;
        # see below
        result[0] = e.exit_code
    invoker_ua.stop()

    # Write any error to stderr, otherwise write hook result to stdout
    error = error_file.getvalue().rstrip()
    if error:
        error_printed = False
        # Unless verbose, ignore the traceback produced by hook commands
        if not options.verbose:
            lines = error.split("\n")
            if lines[0] == "Traceback (most recent call last):":
                lines.pop(0)
                print >> sys.stderr, lines[0]
                error_printed = True
            # Omit wrapper script from error
            prefix = temp_dir + "/jitsu-run-as-hook-script: 2: "
            if lines[0].startswith(prefix):
                print >> sys.stderr, lines[0][len(prefix):]
                error_printed = True
        if not error_printed:
            print >> sys.stderr, error
        if result[0] == 0:
            result[0] = 1  # Return a nonzero status code result
    else:
        value = log_file.getvalue().rstrip()
        if value:
            print value


if __name__ == '__main__':
    main()
