#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  anytop.py
#  anytop
#
#  Created by Lars Yencken on 2011-09-22.
#  Copyright 2011 Lars Yencken. All rights reserved.
#

"""
Live updating frequency distributions on streaming data.
"""

import os
import sys
import optparse
import curses
import time
import threading
import heapq
import re
import logging
from collections import defaultdict, deque
import codecs
import locale

locale.setlocale(locale.LC_ALL, 'en_US.UTF-8')

COLOR_PATTERN = re.compile("\x1b\[[0-9]*(;[0-9]*)?m", re.UNICODE)

MAX_IO_RETRIES = 2

def anytop(win, istream=sys.stdin):
    "Visualize the incoming lines by their distribution."
    istream = codecs.getreader('utf8')(istream)
    _init_win(win)
    dist = defaultdict(int)
    lock = threading.Lock()
    logging.debug('Starting UI thread')
    ui = AnyTopUI(win, lock, dist=dist)
    ui.start()

    color_pattern = re.compile("\x1b\[[0-9]*(;[0-9]*)?m", re.UNICODE)

    try:
        for line in _robust_line_iter(istream):
            if ui.error:
                return ui.error
            # trim and remove shell colors
            key = color_pattern.sub(u'', line.rstrip())

            logging.debug('INPUT: requesting lock')
            lock.acquire()
            logging.debug('INPUT: lock acquired')
            dist[key] += 1
            lock.release()
            logging.debug('INPUT: lock released')

        logging.debug('INPUT: finished')

        # wait for CTRL-C
        # XXX we should display that the input was exhausted
        logging.debug('INPUT: exhausted')
        while True:
            time.sleep(3600)

    except KeyboardInterrupt:
        logging.debug('INPUT: got CTRL-C')
        ui.stop()
        ui.join()

    finally:
        ui.stop()
        ui.join()

    return ui.error

def _robust_line_iter(istream):
    "Read lines, continuing if a blocking system call gets interrupted."
    max_retries = MAX_IO_RETRIES
    tries_since_success = 0

    lines = iter(istream)
    while True:
        try:
            l = lines.next()
            tries_since_success = 0
            yield l

        except StopIteration:
            break

        except IOError:
            # sometimes we can get interrupted on a blocking read
            # let's retry a few times
            tries_since_success += 1

            if tries_since_success >= MAX_IO_RETRIES:
                # something's really wrong
                break

def anytop_window(win, istream=sys.stdin, n=200):
    "Visualize the incoming lines by their distribution."
    istream = codecs.getreader('utf8')(istream)
    _init_win(win)
    logging.debug('Starting UI thread %s' % repr(n))
    queue = deque([], n)
    lock = threading.Lock()
    ui = AnyTopUI(win, lock, queue=queue, dist=None)
    ui.start()
    try:
        for line in _robust_line_iter(istream):
            if ui.error:
                return ui.error
            key = COLOR_PATTERN.sub(u'', line.rstrip())
            logging.debug('INPUT: requesting lock')
            lock.acquire()
            logging.debug('INPUT: lock acquired')
            logging.debug('INPUT: queue size %s' % len(queue))

            queue.append(key)
            lock.release()

        logging.debug('INPUT: exhausted')
        while True:
            time.sleep(3600)

    except KeyboardInterrupt:
        logging.debug('INPUT: got CTRL-C')
        ui.stop()
        ui.join()

    finally:
        ui.stop()
        ui.join()

    return ui.error

def _init_win(win):
    curses.start_color()
    curses.use_default_colors()
    curses.curs_set(0)
    win.nodelay(1)

class AnyTopUI(threading.Thread):
    'The ncurses user interface thread.'
    def __init__(self, win, lock, dist=None, queue=None):
        self.win = win
        self.dist = dist
        self.queue = queue
        self.lock = lock
        self.error = None
        self._stop = threading.Event()
        super(AnyTopUI, self).__init__()

    def stop(self):
        logging.debug('UI: flagged as stopped')
        self._stop.set()

    def stopped(self):
        return self._stop.is_set()

    def run(self):
        while True:
            if self.stopped():
                return
            logging.debug('UI: requesting lock')
            self.lock.acquire()
            logging.debug('UI: lock acquired')
            try:
                self.refresh_display()
            except Exception, e:
                self.lock.release()
                self.error = e
                return
            self.lock.release()
            logging.debug('UI: lock released')
            time.sleep(1)

    def refresh_display(self):
        "Redraw the screen with the current data."
        logging.debug('UI: refreshing the display')
        self.win.erase()
        height, width = self.win.getmaxyx()
        logging.debug('UI: size is %d x %d' % (width, height))
        if self.dist is None:
            d = defaultdict(int)
            for t in self.queue:
                d[t] += 1
        else:
            d = self.dist

        n = len(d)
        s = sum(d.itervalues())

        largest_keys = heapq.nlargest(min(n, height - 2), d, key=d.__getitem__)
        largest = [(d[l], l) for l in largest_keys]

        logging.debug('UI: redraw call')
        self.win.redrawwin()
        self.win.addstr(0, 0, '%d keys, %d counts' % (n, s))
        if largest:
            w = max(6, len(str(max(c for (c, k) in largest))))
            k_len = width - w - 4
            template = u' %%%dd  %%s' % w
            for i, (c, k) in enumerate(largest):
                line = (template % (c, k[:k_len]))[:width]
                try:
                    self.win.addstr(i + 2, 0, line.encode('utf8'))
                except:
                    raise Exception("couldn't draw: '%s'"
                    % line.encode('utf8'))

        logging.debug('UI: refresh')
        self.win.refresh()

#----------------------------------------------------------------------------#

def _create_option_parser():
    usage = \
"""%prog [options]

Live updating frequency distributions on streaming data. Like top, but for any
line-by-line input."""

    parser = optparse.OptionParser(usage)
    parser.add_option('-l', action='store', dest='window', type='int',
            help='Only shows stats on rolling window of n lines.')
    parser.add_option('--debug', action='store_true', dest='debug',
            help='Enable debug logging.')

    return parser

def main(argv):
    parser = _create_option_parser()
    (options, args) = parser.parse_args(argv)

    if args:
        parser.print_help()
        sys.exit(1)

    if options.debug:
        if os.path.exists('debug.log'):
            os.remove('debug.log')

        logging.basicConfig(filename='debug.log', level=logging.DEBUG)

    if options.window:
        err = curses.wrapper(anytop_window, n=options.window)
    else:
        err = curses.wrapper(anytop)

    if err:
        raise err

#----------------------------------------------------------------------------#

if __name__ == '__main__':
    main(sys.argv[1:])

