#!/usr/bin/env python
"""
NetLogger Machine Info daemon
"""
__rcsid__ = "$Id$"
__author__ = "Dan Gunter <dkgunter@lbl.gov>"

# Std library imports
import glob
import logging
import optparse
import os
import pprint
import signal
import socket
import sys
import time

# Third-party imports
import bson
if hasattr(bson, 'dumps'):
    # standalone bson
    bson_encode = bson.dumps
else:
    # pymongo's bson
    bson_encode = bson.BSON.encode

# Local imports
from netlogger import nlmi
from netlogger import util
from netlogger.nllog import OptionParser, get_logger, get_root_logger
from netlogger import xsplib

XSP_TYPE = 0x20 # xsp msg type

#
# Signal handlers
#

# Stop things that are in a loop
g_stop = False

def on_kill(signo, frame):
    """Signal handler for a graceful exit.
    """
    global g_stop
    log = get_logger(__file__)
    log.warn("killed", signo=signo)
    g_stop = True

def on_hup(signo, frame):    
    return

def run(xsp_sess, iterations=None, samples_per_iter=None, dt=None,
        param={}):
    """Main loop.
    """
    assert iterations is not None, "iter not defined"
    assert samples_per_iter is not None, "samples_per_iter not defined"
    assert dt is not None, "dt not defined"

    debug_mode =  xsp_sess is None
    log = get_logger(__file__ + ".run")
    log.info("start")
    iface_meta = { }
    nd_type = ':'.join((nlmi.BASE_ET, 'net', 'dev'))
    st_type = ':'.join((nlmi.BASE_ET, 'stat'))

    # /proc/net/dev initialization. Set interface metadata
    nd_stats = nlmi.get_net_dev(**param)
    b = nlmi.Block()
    t0 = time.time()
    iface_meta = nlmi.build_iface_meta(b, nd_stats, ts=t0, dt=dt)
    # /proc/stat initializaton.
    smeta = nlmi.MetaBlock(event_type='host', subject=dict(hostname=nlmi.get_hostname()))
    b.meta.append(smeta)
    host_info = nlmi.HostInfo(**param)
    # Loop initialization
    sampnum = 0
    # function that will return True when sampling is done
    _complete = lambda num: g_stop or (iterations > 0 and num > iterations*samples_per_iter)
    while not _complete(sampnum):
        nd_data = { } # /proc/net/dev, interface => DataBlock
        samp_start = sampnum
        st_data_arr = [ ]
        for j in xrange(samples_per_iter):
            log.debug("sample.start", iter=sampnum, loop=j)
            # /proc/net/dev
            _ts = time.time()
            nd_stats = nlmi.get_net_dev(**param)
            # Loop over each interface.
            for iface, imeta in iface_meta.iteritems():
                if not iface in nd_data:
                    nd_data[iface] = nlmi.DataBlock(nd_type, meta=imeta)
                nd_data[iface].add_named_values(_ts, nd_stats[iface])
            # /proc/stat
            t, i = time.time(), 0
            # Loop over top-level keys in host info, i.e. the CPUs.
            for info_key, info_value in host_info.get_info().items():
                etype = st_type + '.' + info_key
                if j == 0:
                    # First time, create new DataBlock for i-th metric
                    # and append it to the array (as i-th data block).
                    st_data = nlmi.DataBlock(etype, meta=smeta)
                    st_data.add_named_values(t, info_value)
                    st_data_arr.append(st_data)
                else:
                    # Subsequent times, add i-th metric values to the
                    # existing i-th data block.
                    st_data_arr[i].add_named_values(t, info_value)
                    i += 1
            # Increment and loop
            sampnum += 1
            log.debug("sample.end", iter=sampnum, loop=j, status=0)
            if _complete(sampnum): break
            time.sleep(dt)
        # add samples to main block
        samp_range = (samp_start, sampnum - 1)
        for nd in nd_data.values():
            nd.set_sample_range(*samp_range)
            b.data.append(nd)
        for st_data in st_data_arr:
            st_data.set_sample_range(*samp_range)
            b.data.append(st_data)
        # send message to xspd
        if debug_mode:
            print("\nDATA\n----")
            print(pprint.pformat(b.as_dict()))
        else:
            data = bson_encode(b.as_dict())
            xsp_sess.send_msg(data, len(data), XSP_TYPE)
        b.clear()#_data()
    log.info("end", status=0)
    return 0

def main():
    """Program entry point.
    """
    global g_stop
    g_stop = False
    status = 0
    
    usage = "%prog [options]"
    desc = ' '.join(__doc__.split())
    parser = OptionParser(usage=usage, description=desc, can_be_daemon=True)
    parser.add_option("--debug", dest="dbg", action="store_true", default=False,
            help="Debug mode: dump to stdout instead of connecting to a host")
    parser.add_option("--host", dest="host", default="localhost",
                      help="XSP receiver host (default=%default)")
    parser.add_option("--port", dest="port", action="store", type="int", default=5006,
                      help="XSP receiver port (default=%default)")
    parser.add_option("--dir", dest="proc", action="store", default="/proc",
                      help="Root of Linux 'proc' files (default=%default)")
    group = optparse.OptionGroup(parser, "Sampling options")
    group.add_option("--dt", dest="dt", type="float", default=1.0,
                     help="Sampling interval in seconds (default=%default)")
    group.add_option("--iter", dest="iter", type="int", default=0,
                     help="Number of iterations, 0=forever (default=%default). "
                     "One message is sent to XSP for each iteration.")
    group.add_option("--samples", dest="samples", type="int", default=1,
                     help="Number of samples per iteration (default=%default)")
    parser.add_option_group(group)
    options, args = parser.parse_args(sys.argv[1:])

    # Parse args
    log = get_logger(__file__)  # Must come after parsing args
    log.debug("parse.args.start")
    log.debug("parse.args.end", status=0)

    # Set up signal handlers
    log.debug("init.signals.start")
    util.handleSignals(
        (on_kill, ('SIGTERM', 'SIGINT', 'SIGUSR2')),
        (on_hup, ('SIGHUP',)) )
    log.debug("init.signals.end", status=0)

    if options.dbg:
        if options.daemon:
            parser.error("Cannot use daemon mode with debug mode")
            return(-2)
        xsp_sess = None
    else:
        log.debug("connect.start")
        try:
            xsp_sess = xsplib.XSPSession()
            xsp_sess.connect(options.host, options.port)
        except Exception, err:
            log.critical("connect.error", msg="failed connection to xsp",
                         host=options.host, port=options.port, details=util.traceback())
            return(-1)
        log.debug("connect.end", status=0)


    # If daemon, daemonize
    if options.daemon:
        rootlog = get_root_logger()
        log.debug("daemonize.start")
        try:
            util.daemonize(log=log, root_log=rootlog, close_fds=False)
        except Exception, err:
            log.critical("daemonize.error")
            try:
                xsp_sess.close()
            except socket.error:
                pass
            status = -1
        log.debug("daemonize.end", status=status)

    # Run the probe
    if status == 0:
        try:
            status = run(xsp_sess, iterations=options.iter,
                         samples_per_iter=options.samples, dt=options.dt,
                         param={'dirname':options.proc})
        except Exception, err:
            try:
                if xsp_sess:
                    xsp_sess.close()
            except socket.error:
                pass
            log.critical("probe.error", msg="error in main loop", details=util.traceback())
            status = -2

    if status == 0:
        if xsp_sess:
            xsp_sess.close()

    log.info("run.end", status=status)
    return(status)

if __name__ == '__main__':
    sys.exit(main())
