#!/usr/bin/env python

"""
   fpoll.py - a tool for polling RTYServer clusters
   May be used manually or in automations.

   Wiki: https://wiki.yandex-team.ru/users/yrum/fpoll
"""


import re
import sys
import math
import Queue
import urllib2
import argparse

from threading import Thread, Lock
from socket import timeout as SocketTimeout

W_COUNT=40

class FetchError(Exception):
    pass


class Opts(object):
    def __init__(self):
        self.commands = []
        self.require_success = False  # Stop if command failed
        self.verbose_fails = False
        self.inverse_match = False
        self.re = None
        self.timeout = 5
        self.port_delta = 3


class AsyncTask(object):
    def __init__(self, host, port, opts):
        assert opts is not None
        self.opts = opts
        self.iname = "%s:%d" % (host, port)
        self.ctlUrl = "http://%s:%d/" % (host, port + opts.port_delta)
        self.success = False
        self.exception = False
        self.result = ""

    def _fetchUrl(self, url):
        retryCnt=1
        while retryCnt >= 0:
            try:
                con = urllib2.urlopen(url, timeout=opts.timeout)
                return con.read()
            except urllib2.URLError as e:
                if retryCnt == 0:
                    return ("> request error : %s" % (e.reason,),)
            except urllib2.HTTPError as e:
                if retryCnt == 0:
                    return ("> HTTP error : %d %s" % (e.code, e.reason),)
            except SocketTimeout:
                return ("> HTTP timeout",)
            except:
                pass
            retryCnt -= 1

        return None

    def run(self):
        ok = False;
        for command in self.opts.commands:
            url = self.ctlUrl + command
            result = self._fetchUrl(url)
            if result is not None:
                if not isinstance(result, basestring):
                    result = result[0] if opts.verbose_fails else None # a error
                    # "Hard" failure: opts.require_success and any HTTPError in the row of requests -> stop it
                    self.exception = True
                    ok = False
                elif opts.re is not None:
                    # "Soft" failure: opts.require_success and first request matches -> do the rest
                    ok = ok or (bool(opts.re.search(result)) ^ bool(opts.inverse_match))
                else:
                    ok = True

                if result is not None:
                    self.result += result.rstrip() + "\n"

            # if opts.require_success is not set, will not break whatever happens
            if not ok and opts.require_success:
                break
        self.success = ok

    def __repr__(self):
        return "<%s %s at %s>" % (self.__class__.__name__, self.result, id(self))

    def __str__(self):
        return self.instance


class Worker(Thread):
    def __init__(self, taskQueue, lock):
        super(Worker, self).__init__()
        self.cancelled = False

        self.q = taskQueue
        self.lock = lock

    def run(self):
        while True:
            try:
                task = self.q.get(True,1)
            except Queue.Empty:
                if self.cancelled:
                    break
                continue
            task.run()
            self.q.task_done()

    def cancel(self):
        self.cancelled = True

def startWorkers(queue):
    workers = []
    lock = Lock()
    for i in xrange(W_COUNT):
        w = Worker(queue, lock)
        w.start()
        workers.append(w)
    return workers

def shutdownWorkers(workers):
    map(lambda w: w.cancel(), workers)


def splitInstance(i):
    return (i.split(":")[0], int(i.split(":")[1]))

def create_tasks(tasks, inp):
    for line in inp.readlines():
        iname = line.strip()
        if len(iname) == 0:
            continue

        pos = iname.find('@') # an automatic soltuion for sky listinstances without --no_tag parameter
        if pos != -1:
            iname = iname[:pos]

        host, port = splitInstance(iname)
        tasks.append(AsyncTask(host, port, opts))

def parse_tass(tass):
    # Preprocess json
    lines = tass.replace(',["',',\n["').splitlines()
    chars_ignored = ',[]{}'
    return '\n'.join(sorted([ s.strip(chars_ignored) for s in lines ]))

if __name__ == "__main__":
    task_queue = Queue.Queue()
    parser = argparse.ArgumentParser()

    parser.add_argument("commands", nargs="+", help="One or more commands to execute sequentially.\nThe first one may be used as a condition\n(execution stops for that instance when the regex fails)")
    parser.add_argument("-a", "--all", dest="print_all", action="store_true", help="report all instances, including failures")
    parser.add_argument("-e", "--stop-at-error", action="store_true", help="execute next command only if the previous is ok")
    parser.add_argument("-E", "--verbose-fails", action="store_true", help="print error messages on HTTP errors")
    parser.add_argument("-n", dest="print_names", action="count", help="annotate with instance names (hosts:port).\nUse twice to annotate each string")
    parser.add_argument("-d", dest="delimiter", help="delimiter between results (\\n for empty line) ")
    parser.add_argument("-r", dest="regex", help="regex checker for the response")
    parser.add_argument("-v", dest="inverse", action="store_true", help="inverse match")
    parser.add_argument("-f", dest="filename", help="read instances from FILENAME instead of stdin")
    parser.add_argument("-t", dest="timeout", type=int, default=5, help="timeout for each command")
    parser_adv = parser.add_argument_group("advanced arguments")
    parser_adv.add_argument("--port-delta", metavar='N', type=int, default=3, help="override port_delta") # 0 would mean 'ask the search port'
    parser_adv.add_argument("--json", action="store_true", help="force json processing")
    parser_adv.add_argument("--no-json", action="store_true", help=argparse.SUPPRESS)
    parser_adv.add_argument("--debug", action="store_true", help=argparse.SUPPRESS) # use single thread

    args = parser.parse_args()

    if not args.json and not args.no_json and len(args.commands) == 1 and args.commands[0] == "tass":
        args.json = True

    opts = Opts()
    opts.commands = args.commands
    opts.require_success = args.stop_at_error
    opts.inverse_match = args.inverse
    opts.verbose_fails = args.verbose_fails
    opts.port_delta = args.port_delta
    if args.regex is not None:
        opts.re = re.compile(args.regex)
    opts.timeout = args.timeout

    tasks = []
    if args.filename:
        with open(args.filename, 'r') as inp:
            create_tasks(tasks, inp)
    else:
        create_tasks(tasks, sys.stdin)

    multithreaded = (not args.debug) and len(tasks) > 1

    if multithreaded:
        workers = startWorkers(task_queue)
        for task in tasks:
            task_queue.put(task)

        task_queue.join() # wait till all workers finish their jobs
        shutdownWorkers(workers)
    else:
        for task in tasks:
            task.run()

    #
    # Print results
    #
    delimiter=args.delimiter;
    if delimiter:
        delimiter = delimiter.decode('string_escape')

    for task in tasks:
        if not task.success and not args.print_all and not (task.exception and args.verbose_fails):
            continue
        if args.delimiter:
            sys.stdout.write(delimiter)

        iname_prefix = "%s\t" % task.iname
        if args.print_names == 1:
            sys.stdout.write(iname_prefix)

        res = task.result.strip();
        if args.json:
            res = parse_tass(res)

        if args.print_names > 1:
            lines = []
            for line in res.splitlines():
                lines.append(iname_prefix + line)
            if len(lines) == 0:
                lines.append(iname_prefix); # force empty line if no data and -nn
            res = "\n".join(lines)

        print res

