from twisted.runner.procmon import ProcessMonitor
from twisted.python import log, reflect
import twisted.protocols.basic
from twisted.internet import reactor, protocol
from optparse import OptionParser, SUPPRESS_HELP
import sys, os, time, subprocess
import pickle
import logging

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s [%(name)s] %(levelname)s: %(message)s')

processes = {'lrsscanner': 'youtomb.scan.lrsscanner',
             'deadscanner': 'youtomb.scan.deadscanner',
             'toptodayexplorer': 'youtomb.explore.toptodayexplorer',
             'mosttodayexplorer': 'youtomb.explore.mosttodayexplorer',
             'diggexplorer': 'youtomb.explore.diggexplorer',
             'technoratiexplorer': 'youtomb.explore.technoratiexplorer'}
             
# From twisted.runner.procmon, modified to use Python's logging
class LineLogger(twisted.protocols.basic.LineReceiver):
    tag = None
    delimiter = '\n'
    logger = None
    level = logging.INFO
    
    def lineReceived(self, line):
        if self.logger is None:
            self.logger = logging.getLogger(self.tag)
        self.logger.log(self.level, line)
        
class LoggingProtocol(protocol.ProcessProtocol):
    service = None
    name = None
    empty = 1
    errEmpty = 1

    def connectionMade(self):
        self.output = LineLogger()
        self.output.tag = self.name
        self.output.makeConnection(twisted.runner.procmon.transport)
        self.errOutput = LineLogger()
        self.errOutput.level = logging.ERROR
        self.errOutput.tag = self.name
        self.errOutput.makeConnection(twisted.runner.procmon.transport)

    def outReceived(self, data):
        self.output.dataReceived(data)
        self.empty = data[-1] == '\n'

    def errReceived(self, data):
        self.errOutput.dataReceived(data)
        self.errEmpty = data[-1] == '\n'

    def processEnded(self, reason):
        if not self.empty:
            self.output.dataReceived('\n')
        if not self.errEmpty:
            self.errOutput.dataReceived('\n')
        self.service.connectionLost(self.name)

twisted.runner.procmon.LineLogger = LineLogger
twisted.runner.procmon.LoggingProtocol = LoggingProtocol

def startProcess(self, name):
    if self.protocols.has_key(name):
        return
    p = self.protocols[name] = LoggingProtocol()
    p.service = self
    p.name = name
    args, uid, gid = self.processes[name]
    self.timeStarted[name] = time.time()
    reactor.spawnProcess(p, args[0], args, env=os.environ, uid=uid, gid=gid)
twisted.runner.procmon.ProcessMonitor.startProcess = startProcess

def textFromEventDict(eventDict):
    edm = eventDict['message'] 
    if not edm:
        if eventDict['isError'] and 'failure' in eventDict: 
            text = ((eventDict.get('why') or 'Unhandled Error') 
                    + '\n' + eventDict['failure'].getTraceback()) 
        elif 'format' in eventDict: 
            try:
                text = eventDict['format'] % eventDict
            except:
                text = 'UNFORMATTABLE OBJECT WRITTEN TO LOG with fmt %r, MESSAGE LOST' % (fmtString,)
        else:
            # we don't know how to log this 
            return 
    else: 
        text = ' '.join(map(reflect.safe_str, edm)) 
    return text 

def stdoutEmit(eventDict):
    logging.info(textFromEventDict(eventDict))

def kinit(keytab, principal):
    subprocess.check_call(("/usr/kerberos/bin/kinit","-V","-k","-t",
                           keytab, principal))

def aklog(cells=("athena.mit.edu", "sipb.mit.edu")):
    subprocess.check_call(("/usr/bin/aklog",)+cells)

def main():
    log.addObserver(stdoutEmit)
    logging.info("Starting ProcMon")

    parser = OptionParser()
    parser.add_option("-t", "--threshold", dest="threshold",
                      default=5, type="int",
                      help="number of seconds process has to live before it is considered to be alive", metavar="SECONDS")
    parser.add_option("--with-python", dest="interpreter",
                      default=sys.executable,
                      help="python interpreter to use for launching processes",
                      metavar="PATH")
    parser.add_option("-p", "--principal", dest="principal",
                      default="daemon/youtomb.mit.edu@ATHENA.MIT.EDU",
                      help="Kerberos principal to access files with",
                      metavar="PRINC")
    parser.add_option("--new-pag", dest="new_pag",
                      default=False, action="store_true",
                      help="create a new PAG and kinit before executing")
    parser.add_option("-k", "--keytab", dest="keytab",
                      default="/mit/freeculture/youtomb/youtomb.keytab",
                      help="keytab to use for kinit",
                      metavar="PATH")
    parser.add_option("--pickled-options", dest="pickled_options",
                      help=SUPPRESS_HELP)
    (options, args) = parser.parse_args()
    
    if options.pickled_options:
        (options, args) = pickle.loads(options.pickled_options)
        if options.new_pag:
            aklog()
            def doKrenew():
                kinit(options.keytab, options.principal)
                aklog()
                reactor.callLater(60*60, doKrenew)
            reactor.callLater(60*60, doKrenew)
        else:
            print "Uhh, pickled options but not new_pag?"
            print "Got options:", options
    elif options.new_pag:
        opts = pickle.dumps((options, args))
        os.environ["KRB5CCNAME"] = os.tempnam(None, "yt.krb5cc")
        print "Using KRB5CCNAME=%s" % os.environ["KRB5CCNAME"]
        kinit(options.keytab, options.principal)
        # Need to aklog here so python can read the scripts
        command = ("pagsh", "pagsh", "-c", 'aklog athena sipb; "$@"', "python", options.interpreter, sys.argv[0], "--pickled-options", opts)
        print "Executing", command
        os.execlp(*command)
    
    mon = ProcessMonitor()
    mon.threshold = options.threshold
    for module in args:
        module = processes.get(module, module)
        mon.addProcess(module, [options.interpreter, "-m", module])
    mon.startService()
    reactor.addSystemEventTrigger('before', 'shutdown', mon.stopService)
    reactor.run()

if __name__ == "__main__":
    main()