
# standard
import cStringIO
import optparse
import os
import resource
import signal
import string
import sys
import time
import traceback
from collections import defaultdict

# local
import sink
from core import __version__

# vendor
import gevent, gevent.socket
socket = gevent.socket


# constants
INTERVAL = 10.0
PERCENT = 90.0
MAX_PACKET = 2048

DESCRIPTION = '''
A statsd service in Python + gevent.
'''
EPILOG = '''

'''

# table to remove invalid characters from keys
ALL_ASCII = set(chr(c) for c in range(256))
KEY_VALID = string.ascii_letters + string.digits + '_-.'
KEY_TABLE = string.maketrans(KEY_VALID + '/', KEY_VALID + '_')
KEY_DELETIONS = ''.join(ALL_ASCII.difference(KEY_VALID + '/'))

# error messages
E_BADADDR = 'invalid bind address specified %r'
E_NOSINKS = 'you must specify at least one stats sink'


class Stats(object):

    def __init__(self):
        self.timers = defaultdict(list)
        self.counts = defaultdict(float)
        self.percent = PERCENT
        self.interval = INTERVAL


def daemonize(umask=0027):
    if gevent.fork():
        os._exit(0)
    os.setsid()
    if gevent.fork():
        os._exit(0)
    os.umask(umask)
    fd_limit = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
    if fd_limit == resource.RLIM_INFINITY:
        fd_limit = 1024
    for fd in xrange(0, fd_limit):
        try:
            os.close(fd)
        except:
            pass
    os.open(os.devnull, os.O_RDWR)
    os.dup2(0, 1)
    os.dup2(0, 2)
    gevent.reinit()


def parse_addr(text):
    "Parse a 1- to 3-part address spec."
    if text:
        parts = text.split(':')
        length = len(parts)
        if length== 3:
            return parts[0], parts[1], int(parts[2])
        elif length == 2:
            return None, parts[0], int(parts[1])
        elif length == 1:
            return None, '', int(parts[0])
    return None, None, None


class StatsDaemon(object):

    """
    A statsd service implementation in Python + gevent.
    """

    def __init__(self, bindaddr, sinkspecs, interval, percent, debug=0):
        _, host, port = parse_addr(bindaddr)
        if port is None:
            self.exit(E_BADADDR % bindaddr)
        self._bindaddr = (host, port)

        # TODO: generalize to support more than one sink type.  currently
        # only the graphite backend is present, but we may want to write
        # stats to hbase, redis, etc. - ph

        # construct the sink and add hosts to it
        if not sinkspecs:
            self.exit(E_NOSINKS)
        self._sink = sink.GraphiteSink()
        errors = []
        for spec in sinkspecs:
            try:
                self._sink.add(spec)
            except ValueError, ex:
                errors.append(ex)
        if errors:
            for err in errors:
                self.error(str(err))
            self.exit('exiting.')

        self._percent = float(percent)
        self._interval = float(interval)
        self._debug = debug
        self._sock = None
        self._flush_task = None

        self._reset_stats()

    def _reset_stats(self):
        self._stats = Stats()
        self._stats.percent = self._percent
        self._stats.interval = self._interval

    def exit(self, msg, code=1):
        self.error(msg)
        sys.exit(code)

    def error(self, msg):
        sys.stderr.write(msg + '\n')

    def start(self):
        "Start the service"
        # register signals
        gevent.signal(signal.SIGINT, self._shutdown)

        # spawn the flush trigger
        def _flush_impl():
            while 1:
                gevent.sleep(self._stats.interval)

                # rotate stats
                stats = self._stats
                self._reset_stats()

                # send the stats to the sink which in turn broadcasts
                # the stats packet to one or more hosts.
                try:
                    self._sink.send(stats)
                except Exception, ex:
                    trace = traceback.format_tb(sys.exc_info()[-1])
                    self.error(''.join(trace))

        self._flush_task = gevent.spawn(_flush_impl)

        # start accepting connections
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, 
            socket.IPPROTO_UDP)
        self._sock.bind(self._bindaddr)
        while 1:
            try:
                self._process(*self._sock.recvfrom(MAX_PACKET))
            except Exception, ex:
                self.error(str(ex))

    def _shutdown(self):
        "Shutdown the server"
        self.exit("service exiting", code=0)

    def _process(self, data, _):
        "Process a single packet and update the internal tables."
        parts = data.split(':')
        if self._debug:
            self.error('packet: %r' % data)
        if not parts:
            return

        # interpret the packet and update stats
        stats = self._stats
        key = parts[0].translate(KEY_TABLE, KEY_DELETIONS)
        for part in parts[1:]:
            srate = 1.0
            fields = part.split('|')
            length = len(fields)
            if length < 2:
                continue
            value = fields[0]
            stype = fields[1].strip()

            # timer (milliseconds)
            if stype == 'ms':
                stats.timers[key].append(float(value if value else 0))

            # counter with optional sample rate
            elif stype == 'c':
                if length == 3 and fields[2].startswith('@'):
                    srate = float(fields[2][1:])
                value = float(value if value else 1) * (1 / srate)
                stats.counts[key] += value


def main():
    opts = optparse.OptionParser(description=DESCRIPTION, version=__version__,
        add_help_option=False)
    opts.add_option('-b', '--bind', dest='bind_addr', default=':8125', 
        help="bind [host]:port (host defaults to '')")
    opts.add_option('-s', '--sink', dest='sink', action='append', default=[],
        help="a graphite service to which stats are sent ([host]:port).")
    opts.add_option('-v', dest='verbose', action='count', default=0,
        help="increase verbosity (currently used for debugging)")
    opts.add_option('-f', '--flush', dest='interval', default=INTERVAL,
        help="flush interval, in seconds (default 10)")
    opts.add_option('-p', '--percent', dest='percent', default=PERCENT,
        help="percent threshold (default 90)")
    opts.add_option('-D', '--daemonize', dest='daemonize', action='store_true',
        help='daemonize the service')
    opts.add_option('-h', '--help', dest='usage', action='store_true')

    (options, args) = opts.parse_args()

    if options.usage:
        # TODO: write epilog. usage is manually output since optparse will
        # wrap the epilog and we want pre-formatted output. - ph
        print(opts.format_help())
        sys.exit()

    if options.daemonize:
        daemonize()

    sd = StatsDaemon(options.bind_addr, options.sink, options.interval,
        options.percent, options.verbose)
    sd.start()
 

if __name__ == '__main__':
    main()

