#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os
import time
import argparse
import redis
from redis.exceptions import ConnectionError
from rq import use_redis, Queue, Worker
from rq.utils import gettermsize, make_colorizer

red = make_colorizer('darkred')
green = make_colorizer('darkgreen')
yellow = make_colorizer('darkyellow')


def pad(s, pad_to_length):
    """Pads the given string to the given length."""
    return ('%-' + '%ds' % pad_to_length) % (s,)

def get_scale(x):
    """Finds the lowest scale where x <= scale."""
    scales = [20, 50, 100, 200, 400, 600, 800, 1000]
    for scale in scales:
        if x <= scale:
            return scale
    return x

def state_symbol(state):
    symbols = {
        'busy': red('busy'),
        'idle': green('idle'),
    }
    try:
        return symbols[state]
    except KeyError:
        return state


def show_queues(args):
    while True:
        if len(args.queues):
            qs = map(Queue, args.queues)
        else:
            qs = Queue.all()

        num_jobs = 0
        termwidth, _ = gettermsize()
        chartwidth = min(20, termwidth - 20)

        max_count = 0
        counts = dict()
        for q in qs:
            count = q.count
            counts[q] = count
            max_count = max(max_count, count)
        scale = get_scale(max_count)
        ratio = chartwidth * 1.0 / scale

        if args.interval:
            os.system('clear')

        for q in qs:
            count = counts[q]
            if not args.raw:
                chart = green('|' + '█' * int(ratio * count))
                line = '%-12s %s %d' % (q.name, chart, count)
            else:
                line = '%-12s %d' % (q.name, count)
            print(line)

            num_jobs += count

        # Print summary when not in raw mode
        if not args.raw:
            print('%d queues, %d jobs total' % (len(qs), num_jobs))

        if args.interval:
            time.sleep(args.interval)
        else:
            break

def show_workers(args):
    while True:
        qs = Queue.all()
        ws = Worker.all()

        if args.interval:
            os.system('clear')

        queues = {qname: [] for qname in qs}
        for w in ws:
            for q in w.queues:
                if not q in queues:
                    queues[q] = []
                queues[q].append(w)

        if args.by_queue:
            max_qname = max(map(lambda q: len(q.name), queues.keys())) if queues else 0
            for q in queues:
                if queues[q]:
                    queues_str = ", ".join(sorted(map(lambda w: '%s (%s)' % (w.name, state_symbol(w.state)), queues[q])))
                else:
                    queues_str = '–'
                print '%s %s' % (pad(q.name + ':', max_qname + 1), queues_str)
        else:
            for w in ws:
                print '%s %s: %s' % (w.name, state_symbol(w.state), ', '.join(w.queue_names()))
        print '%d workers, %d queues' % (len(ws), len(queues))

        if args.interval:
            time.sleep(args.interval)
        else:
            break


def parse_args():
    parser = argparse.ArgumentParser(description='RQ command-line monitor.')
    parser.add_argument('--host', '-H', default='localhost', help='The Redis hostname (default: localhost)')
    parser.add_argument('--port', '-p', type=int, default=6379, help='The Redis portnumber (default: 6379)')
    parser.add_argument('--db', '-d', type=int, default=0, help='The Redis database (default: 0)')
    parser.add_argument('--path', '-P', default='.', help='Specify the import path.')

    parent_parser = argparse.ArgumentParser(add_help=False)
    parent_parser.add_argument('--interval', '-i', metavar='N', type=float, default=0, help='Updates stats every N seconds (default: don\'t poll)')

    subparsers = parser.add_subparsers()

    queues_p = subparsers.add_parser('queues', parents=[parent_parser], help='Show queue info')
    queues_p.add_argument('--raw', '-r', action='store_true', default=False, help='Print only the raw numbers, no bar charts')
    queues_p.add_argument('queues', nargs='*', help='The queues to poll')
    queues_p.set_defaults(func=show_queues)

    workers_p = subparsers.add_parser('workers', parents=[parent_parser], help='Show worker activity')
    workers_p.add_argument('--by-queue', '-Q', dest='by_queue', default=False, action='store_true', help='Shows workers by queue')
    workers_p.set_defaults(func=show_workers)

    return parser.parse_args()


def main():
    args = parse_args()

    if args.path:
        sys.path = args.path.split(':') + sys.path

    # Setup connection to Redis
    redis_conn = redis.Redis(host=args.host, port=args.port, db=args.db)
    use_redis(redis_conn)
    try:
        args.func(args)
    except ConnectionError as e:
        print(e)

if __name__ == '__main__':
    main()