'''
Benchmark the configured DNS service.

Verify that DNS more-or-less works. Measure its speed to see if it'll
be a bottleneck.
'''

import sys
import time
import argparse
import asyncio
import random
import os

import cocrawler.dns as dns
import cocrawler.config as config

ARGS = argparse.ArgumentParser(description='CoCrawler dns benchmark')
ARGS.add_argument('--config', action='append')
ARGS.add_argument('--configfile', action='store')
ARGS.add_argument('--no-confighome', action='store_true')
ARGS.add_argument('--count', type=int, default=1000)
ARGS.add_argument('--expect-not-suitable', action='store_true')

args = ARGS.parse_args()

config.config(args.configfile, args.config, confighome=not args.no_confighome)
max_workers = config.read('Crawl', 'MaxWorkers')

ns = config.read('Fetcher', 'Nameservers')
if not isinstance(ns, list):
    ns = [ns]

#resolver = aiohttp.resolver.AsyncResolver(nameservers=ns)  # Can I pass rotate=True into this?
#connector = aiohttp.connector.TCPConnector(resolver=resolver, family=socket.AF_INET)
#session = aiohttp.ClientSession(connector=connector)
exit_value = 0

dns.setup_resolver(ns)


def create_queue():
    queue = asyncio.Queue()

    # add a fake domain to make sure the dns doesn't send unknown hosts to a search
    # note that mail.foo.com and mx.foo.com don't generally get bogus answers, it's foo.com or www.foo.com that do
    for _ in range(2):
        r = random.Random()
        host = str(r.randrange(1000000000)) + str(r.randrange(1000000000)) + str(r.randrange(1000000000))
        queue.put_nowait(host + '.com')

    # read list of domains to query -- from alexa top million
    head, tail = os.path.split(__file__)
    alexa = os.path.join(head, os.pardir, 'examples', 'top-1k.txt')
    alexa_count = 0

    try:
        with open(alexa, 'r') as f:
            for line in f:
                queue.put_nowait(line.strip())
                alexa_count += 1
                if alexa_count > args.count:
                    break
    except FileNotFoundError:
        # the alexa file wasn't available (it is not in the repo) so just do a few
        for _ in range(args.count):
            queue.put_nowait('www.google.com')
    return queue


async def work():
    while True:
        try:
            host = queue.get_nowait()
        except asyncio.queues.QueueEmpty:
            break

        try:
            result = await dns.query(host, 'A')
        except Exception as e:
            result = None
            print('saw exception', e, 'but ignoring it')

        if result is not None:
            if args.expect_not_suitable:
                print('as expected, this nameserver is not suitable for crawling')
            else:
                print('invalid hostname got a result: your nameserver is not suitable for crawling')
                global exit_value
                exit_value = 1

        queue.task_done()


async def main():
    workers = [asyncio.Task(work(), loop=loop) for _ in range(max_workers)]
    await queue.join()
    for w in workers:
        if not w.done():
            w.cancel()

queue = create_queue()
qsize = queue.qsize()

print('workers:', max_workers)
print('configured nameservers:', ns)

t0 = time.time()

loop = asyncio.get_event_loop()
try:
    loop.run_until_complete(main())
except KeyboardInterrupt:
    sys.stderr.flush()
    sys.stdout.flush()
finally:
    loop.stop()
    loop.run_forever()
    loop.close()
    #session.close()

elapsed = time.time() - t0
if not elapsed:
    elapsed = 1
qps = qsize / elapsed

print('processed {} dns calls in {:.1f} seconds, qps = {:.1f}'.format(qsize, elapsed, qps))

sys.exit(exit_value)

@@ -1562,16 +1562,17 @@
 _nowait(
+(
 host + '
@@ -1576,16 +1576,25 @@
 + '.com'
+, 'fake')
 )%0A%0A    #
@@ -1738,16 +1738,12 @@
 r, '
-examples
+data
 ', '
@@ -1817,24 +1817,92 @@
  'r') as f:%0A
+            print('Using top-1k from Alexa, expect a few failures')%0A
             
@@ -1949,16 +1949,17 @@
 _nowait(
+(
 line.str
@@ -1962,16 +1962,25 @@
 .strip()
+, 'real')
 )%0A      
@@ -2190,16 +2190,92 @@
 o a few%0A
+        print('Cannot find top-1k file, so all queries are www.google.com')%0A
         
@@ -2331,16 +2331,17 @@
 _nowait(
+(
 'www.goo
@@ -2344,24 +2344,33 @@
 .google.com'
+, 'real')
 )%0A    return
@@ -2441,16 +2441,22 @@
     host
+, kind
  = queue
@@ -2648,16 +2648,55 @@
  = None%0A
+            if kind is not 'fake':%0A    
         
@@ -2774,16 +2774,35 @@
 not None
+ and kind is 'fake'
 :%0A      

make this actually work