from functools import wraps
from greplin import scales
import celery.bin.base
import collections
import json
import logging
import sys
import threading
import time
import zbxsend


if sys.version_info < (3,):
    from ConfigParser import ConfigParser
    from cStringIO import StringIO
    import thread
else:
    from configparser import ConfigParser
    from io import StringIO
    import _thread as thread


log = logging.getLogger(__name__)
stats = scales.collection(
    '/celery',
    scales.IntStat('started'),

    scales.IntStat('succeeded'),
    scales.IntStat('failed'),
    scales.IntStat('retried'),

    scales.PmfStat('queuetime'),
    scales.PmfStat('runtime'),
)
stats_queue = type('Stats:queues', (object,), {})()
scales._Stats.initChild(stats_queue, 'queues', '', stats)


def task_handler(fn):
    @wraps(fn)
    def wrapper(self, event):
        self.state.event(event)
        task = self.state.tasks.get(event['uuid'])
        return fn(self, event, task)
    return wrapper


class Receiver(object):

    def __init__(self, app):
        self.app = app

    @task_handler
    def on_task_started(self, event, task):
        # XXX We'd like to maybe differentiate this by queue, but
        # task.routing_key is always None, even though in redis it contains the
        # queue name.
        log.debug('Started %s', task)
        stats.started += 1
        if task.sent:
            stats.queuetime = time.time() - task.sent

    @task_handler
    def on_task_succeeded(self, event, task):
        log.debug('Succeeded %s', task)
        stats.succeeded += 1
        if task is not None and task.runtime is not None:
            stats.runtime = task.runtime

    @task_handler
    def on_task_failed(self, event, task):
        log.debug('Failed %s', task)
        stats.failed += 1
        if task is not None and task.runtime is not None:
            stats.runtime = task.runtime

    @task_handler
    def on_task_retried(self, event, task):
        log.debug('Retried %s', task)
        stats.retried += 1
        if task is not None and task.runtime is not None:
            stats.runtime = task.runtime

    def __call__(self, *args, **kw):
        self.state = self.app.events.State()
        kw.setdefault('wakeup', False)

        with self.app.connection() as connection:
            recv = self.app.events.Receiver(connection, handlers={
                'task-started': self.on_task_started,
                'task-succeeded': self.on_task_succeeded,
                'task-failed': self.on_task_failed,
                'task-retried': self.on_task_retried,
                '*': self.state.event,
            })
            recv.capture(*args, **kw)


class Command(celery.bin.base.Command):

    should_stop = False

    def dump_stats(self):
        while not self.should_stop:
            try:
                data = scales.getStats()['celery']
                log.debug(data)
                metrics = {
                    'celery.task.' + x: data.get(x, 0)
                    for x in ['started', 'succeeded', 'failed', 'retried']
                }
                metrics['celery.task.runtime'] = data.get(
                    'runtime', {}).get('median', -1)
                metrics['celery.task.queuetime'] = data.get(
                    'queuetime', {}).get('median', -1)
                discovery = []
                for key, value in data['queues'].items():
                    metrics['celery.queue[%s]' % key] = value
                    discovery.append({'{#QUEUENAME}': key})
                # See <https://www.zabbix.com/documentation/3.0/manual/
                #   /discovery/low_level_discovery#creating_custom_lld_rules>
                # and <https://github.com/jbfavre/python-protobix/blob/1.0.1
                #   /protobix/datacontainer.py#L53>
                metrics['celery.discover.queues'] = json.dumps(
                    {'data': discovery})
                self._send_to_zabbix(metrics)
                log.debug(
                    'Dump thread going to sleep for %s seconds',
                    self.dump_interval)
                time.sleep(self.dump_interval)
            except Exception:
                log.error(
                    'Uncaught exception, preventing thread from crashing.',
                    exc_info=True)

    def _send_to_zabbix(self, metrics):
        if not (self.zabbix_server and self.zabbix_nodename):
            log.warning('No zabbix connection configured, not sending metrics')
            return
        # Work around bug in zbxsend, they keep the fraction which zabbix
        # then rejects.
        now = int(time.time())
        metrics = [zbxsend.Metric(self.zabbix_nodename, key, value, now)
                   for key, value in metrics.items()]
        log.debug(metrics)
        zbxsend.send_to_zabbix(metrics, self.zabbix_server)

    def check_queue_lengths(self):
        while not self.should_stop:
            try:
                lengths = collections.Counter()

                with self.app.connection() as connection:
                    pipe = connection.channel().client.pipeline(
                        transaction=False)
                    for queue in self.app.conf['task_queues']:
                        if not hasattr(stats_queue, queue.name):
                            setattr(type(stats_queue), queue.name,
                                    scales.Stat(queue.name))
                        # Not claimed by any worker yet
                        pipe.llen(queue.name)
                    # Claimed by worker but not acked/processed yet
                    pipe.hvals('unacked')

                    result = pipe.execute()

                unacked = result.pop()
                for task in unacked:
                    task = json.loads(task.decode('utf-8'))
                    lengths[task[-1]] += 1
                unacked = [[-1] for v in unacked]
                unacked = len([x for x in unacked])

                for llen, queue in zip(result, self.app.conf['task_queues']):
                    lengths[queue.name] += llen

                for queue, length in lengths.items():
                    setattr(stats_queue, queue, length)

                time.sleep(self.queuelength_interval)
            except Exception:
                log.error(
                    'Uncaught exception, preventing thread from crashing.',
                    exc_info=True)

    def run(self, **kw):
        receiver = Receiver(self.app)

        try_interval = 1
        while not self.should_stop:
            try:
                try_interval *= 2
                receiver()
                try_interval = 1
            except (KeyboardInterrupt, SystemExit):
                log.info('Exiting')
                receiver.should_stop = True
                self.should_stop = True
                thread.interrupt_main()
                break
            except Exception as e:
                log.error(
                    'Failed to capture events: "%s", '
                    'trying again in %s seconds.',
                    e, try_interval, exc_info=True)
                time.sleep(try_interval)

    def prepare_args(self, *args, **kw):
        options, args = super(Command, self).prepare_args(*args, **kw)
        self.app.log.setup(
            logging.DEBUG if options.get('verbose') else logging.INFO)

        self._configure_zabbix(options)
        self.dump_interval = options.pop('dump_interval')
        threading.Thread(target=self.dump_stats).start()

        self.queuelength_interval = options.pop('queuelength_interval')
        if self.queuelength_interval:
            threading.Thread(target=self.check_queue_lengths).start()

        return options, args

    def _configure_zabbix(self, options):
        self.zabbix_server = options.pop('zabbix_server', None)
        self.zabbix_nodename = options.pop('zabbix_nodename', None)

        agent_config = options.get('zabbix_agent_config')
        if (agent_config and
                (not self.zabbix_server or not self.zabbix_nodename)):
            log.debug('Using zabbix agent config %s', agent_config)
            text = open(agent_config).read()
            text = '[general]\n' + text
            config = ConfigParser()
            config.readfp(StringIO(text))
            if not self.zabbix_server:
                self.zabbix_server = config.get('general', 'Server')
                if ',' in self.zabbix_server:
                    self.zabbix_server = self.zabbix_server.split(',')[0]
            if not self.zabbix_nodename:
                self.zabbix_nodename = config.get('general', 'Hostname')
        log.debug('Using zabbix server %s', self.zabbix_server)
        log.debug('Using zabbix nodename %s', self.zabbix_nodename)

    def add_arguments(self, parser):
        parser.add_argument(
            '--verbose', help='Enable debug logging', action='store_true')

        parser.add_argument('--zabbix-server', help='Zabbix server name')
        parser.add_argument(
            '--zabbix-nodename', help='Report to zabbix as this hostname')
        parser.add_argument(
            '--zabbix-agent-config', help='Path to zabbix_agentd.conf, '
            'to read zabbix server+port and hostname from.')

        parser.add_argument(
            '--dump-interval',
            help='Send metrics to zabbix every x seconds',
            type=int, default=60)

        parser.add_argument(
            '--queuelength-interval',
            help='Check queue lengths every x seconds (0=disabled)',
            type=int, default=0)
