"""Compare DNS on two optimized uvloop builds without external DNS traffic.

Run each sample in a fresh process. The parent alternates build order and
reports medians; individual samples and build paths remain in JSON output.
"""
import argparse
import asyncio
import concurrent.futures
import json
import os
import pathlib
import platform
import socket
import statistics
import subprocess
import sys
import threading
import time


def child(args):
    sys.path.insert(0, str(pathlib.Path(args.repo).resolve()))
    import uvloop
    import uvloop.loop
    loop = uvloop.new_event_loop()
    assert not loop._debug_cc, 'benchmark requires an optimized build'
    assert pathlib.Path(uvloop.loop.__file__).is_relative_to(
        pathlib.Path(args.repo).resolve())
    asyncio.set_event_loop(loop)

    async def lookup(reverse=False, numeric=False):
        if reverse:
            return await loop.getnameinfo(
                ('127.0.0.1', 443), socket.NI_NUMERICHOST | socket.NI_NUMERICSERV)
        return await loop.getaddrinfo(
            '127.0.0.1' if numeric else 'localhost', 443,
            family=socket.AF_INET, type=socket.SOCK_STREAM)

    async def run():
        if args.case == 'executor_saturated':
            # Explicitly use four workers so this case is reproducible across
            # machines. Occupy every worker with sleeping application work.
            loop.set_default_executor(
                concurrent.futures.ThreadPoolExecutor(max_workers=4))
            release = threading.Event()
            ready = [threading.Event() for _ in range(4)]

            def application_work(event):
                event.set()
                release.wait()

            tasks = [loop.run_in_executor(None, application_work, e)
                     for e in ready]
            while not all(e.is_set() for e in ready):
                await asyncio.sleep(0.001)
            timer = loop.call_later(0.2, release.set)
            start = time.perf_counter()
            try:
                await lookup()
                elapsed = time.perf_counter() - start
            finally:
                release.set()
                timer.cancel()
                await asyncio.gather(*tasks)
            return {'dns_latency_ms': elapsed * 1000}

        reverse = args.case.startswith('reverse')
        numeric = args.case == 'numeric_fast_path'
        concurrency = 32 if args.case.endswith('concurrent') else 1
        count = 32768 if numeric else (4096 if reverse else 8192)
        for _ in range(100):
            await lookup(reverse, numeric)

        async def worker():
            for _ in range(count // concurrency):
                result = await lookup(reverse, numeric)
                assert result

        start = time.perf_counter()
        await asyncio.gather(*(worker() for _ in range(concurrency)))
        elapsed = time.perf_counter() - start
        return {'lookups_per_second': count / elapsed,
                'microseconds_per_lookup': elapsed * 1e6 / count,
                'lookups': count, 'concurrency': concurrency}

    try:
        result = loop.run_until_complete(run())
        loop.run_until_complete(loop.shutdown_default_executor())
    finally:
        loop.close()
    result.update(module=uvloop.loop.__file__, version=uvloop.__version__,
                  libuv_version=uvloop.loop.libuv_get_version())
    print(json.dumps(result))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--repo')
    parser.add_argument('--case')
    parser.add_argument('--baseline')
    parser.add_argument('--candidate')
    parser.add_argument('--samples', type=int, default=7)
    parser.add_argument('--output')
    args = parser.parse_args()
    if args.repo:
        child(args)
        return
    cases = ('forward_sequential', 'forward_concurrent',
             'reverse_sequential', 'reverse_concurrent',
             'numeric_fast_path', 'executor_saturated')
    result = {'environment': {'python': sys.version,
                             'platform': platform.platform(),
                             'cpu_count': os.cpu_count(),
                             'uv_threadpool_size': os.getenv('UV_THREADPOOL_SIZE', 'default (4)'),
                             'samples': args.samples}, 'cases': {}}
    env = dict(os.environ)
    env.pop('PYTHONASYNCIODEBUG', None)
    env.pop('LD_PRELOAD', None)
    env['UV_THREADPOOL_SIZE'] = '4'
    for case in cases:
        samples = {name: [] for name in ('baseline', 'candidate')}
        for i in range(args.samples):
            order = ('baseline', 'candidate') if i % 2 == 0 else ('candidate', 'baseline')
            for name in order:
                output = subprocess.check_output(
                    [sys.executable, str(pathlib.Path(__file__).resolve()),
                     '--repo', getattr(args, name), '--case', case],
                    text=True, env=env, timeout=60)
                samples[name].append(json.loads(output))
        key = 'dns_latency_ms' if case == 'executor_saturated' else 'lookups_per_second'
        medians = {name: statistics.median(s[key] for s in runs)
                   for name, runs in samples.items()}
        result['cases'][case] = {'samples': samples, 'median': medians,
                                 'candidate_over_baseline': medians['candidate'] / medians['baseline']}
        print(case, json.dumps(medians), flush=True)
        pathlib.Path(args.output).write_text(json.dumps(result, indent=2) + '\n')


if __name__ == '__main__':
    main()