"""
This script takes a sitemap file generated by `spider.py` and
spins up a number of workers (threads/processes/greenlets) to
hit those URLs.

A few additional features would be helpful here like:
 - remembering stats about the URLs requested
   + min, max, and avg response times
   + min, max, and avg response sizes
 - failure rates for URLs
 - response code counts (e.g.: 95 2xx, 4 4xx, and 1 5xx)

This script runs forever (or until it's killed) checking
these URLs and periodically outputing stats.

"""
import collections
from multiprocessing import Process
import logging
import signal
import time
import timeit

import gevent
from gevent import pool
from gevent import monkey
import requests
import yaml


monkey.patch_socket()
requests_log = logging.getLogger("requests")


def sitemap_to_urls(sitemap_path):
    with open(sitemap_path, "r") as sitemap_file:
        sitemap = yaml.load(sitemap_file.read())

    urls = []

    def append_paths(parent_path, d, urls):
        for key in d:
            url_path = parent_path + (key,)
            url = "/".join(url_path)
            urls.append(url)
            append_paths(url_path, d[key], urls)
    append_paths((), sitemap, urls)

    return urls



def takeN(size, iterable):
    batch = []
    for n, i in enumerate(iterable):
        batch.append(i)
        if n % size == 0:
            yield batch
            batch = []
        yield batch


base = ""


# STATS!
START_TIME = time.time()
TOTAL_ELAPSED = time.time() - START_TIME
LAST_NOTIFIED = 0.0
NOTIFICATION_INTERVAL = 60.0
COUNTER = collections.Counter()


def get_url(url):
    global LAST_NOTIFIED, TOTAL_ELAPSED

    absolute_url = base + url
    logging.debug("getting: %s ... ", absolute_url)
    start = timeit.default_timer()

    try:
        resp = requests.get(absolute_url, timeout=30.0)
    except requests.exceptions.Timeout:
        stop = timeit.default_timer()
        elapsed = stop - start
        logging.debug("%s: Timed out (%s)", absolute_url, elapsed)
        key = "{}:TIMEOUT".format(absolute_url)
        COUNTER.update({key: 1})
    except requests.exceptions.SSLError:
        stop = timeit.default_timer()
        elapsed = stop - start
        logging.debug("%s: SSL Error (%s)", absolute_url, elapsed)
        key = "{}:SSL_ERROR".format(absolute_url)
        COUNTER.update({key: 1})
    except requests.exceptions.ConnectionError:
        stop = timeit.default_timer()
        elapsed = stop - start
        logging.debug("%s: Connection Error (%s)", absolute_url, elapsed)
        key = "{}:CONNECT_ERROR".format(absolute_url)
        COUNTER.update({key: 1})
    else:
        stop = timeit.default_timer()
        elapsed = stop - start
        logging.debug("%s: %s (%s) (%4f)", absolute_url, resp.status_code, resp.reason, elapsed)
        key = "{}:{}".format(absolute_url, resp.status_code)
        COUNTER.update({key: 1})
    finally:
        TOTAL_ELAPSED = time.time() - START_TIME
        if TOTAL_ELAPSED > LAST_NOTIFIED + NOTIFICATION_INTERVAL:
            LAST_NOTIFIED = TOTAL_ELAPSED
            logging.info("most common response stats:")
            for x in COUNTER.most_common(10):
                logging.info("\t%s", x)


def run(inflight, urls, ttl):
    if ttl > 0:
        signal.alarm(ttl)

    gevent.reinit()
    p = pool.Pool(inflight)
    while True:
        p.map(get_url, urls)
        p.join()


def main(arguments):
    global base

    base = arguments.base_url
    sitemap_path = arguments.sitemap_path
    inflight = arguments.inflight
    processes = arguments.processes
    debug = arguments.debug
    ttl = arguments.ttl

    if debug:
        logging.basicConfig(level=logging.DEBUG,
                            format="%(asctime)s - %(levelname)s - %(message)s")
    else:
        requests_log.setLevel(logging.WARNING)
        logging.basicConfig(level=logging.INFO,
                            format="%(asctime)s - %(levelname)s - %(message)s")

    logging.info("spinneret generating traffic - URL=%s N=%s P=%s",
                 base, inflight, processes)

    urls = sitemap_to_urls(sitemap_path)

    if processes == 1:
        run(inflight, urls, ttl)
    else:
        process_list = [Process(target=run, args=(inflight, urls, ttl)) for x in xrange(processes)]
        for proc in process_list:
            proc.start()
        for proc in process_list:
            proc.join()
