'''
This test client runs against the Autobahn test server. It is based on the
test_client.py in wsproto.
'''
import argparse
import json
import logging
import sys

import trio
from trio_websocket import open_websocket_url, ConnectionClosed


AGENT = 'trio-websocket'
MAX_MESSAGE_SIZE = 16 * 1024 * 1024
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('client')


async def get_case_count(url: str) -> int:
    url = url + '/getCaseCount'
    async with open_websocket_url(url) as conn:
        case_count = await conn.get_message()
        logger.info('Case count=%s', case_count)
    return int(case_count)


async def get_case_info(url: str, case: str) -> object:
    url = f'{url}/getCaseInfo?case={case}'
    async with open_websocket_url(url) as conn:
        return json.loads(await conn.get_message())


async def run_case(url: str, case: str) -> None:
    url = f'{url}/runCase?case={case}&agent={AGENT}'
    try:
        async with open_websocket_url(url, max_message_size=MAX_MESSAGE_SIZE) as conn:
            while True:
                data = await conn.get_message()
                await conn.send_message(data)
    except ConnectionClosed:
        pass


async def update_reports(url: str) -> None:
    url = f'{url}/updateReports?agent={AGENT}'
    async with open_websocket_url(url) as conn:
        # This command runs as soon as we connect to it, so we don't need to
        # send any messages.
        pass


async def run_tests(args: argparse.Namespace) -> None:
    logger = logging.getLogger('trio-websocket')
    if args.debug_cases:
        # Don't fetch case count when debugging a subset of test cases. It adds
        # noise to the debug logging.
        case_count = None
        test_cases = args.debug_cases
    else:
        case_count = await get_case_count(args.url)
        test_cases = list(range(1, case_count + 1))
    exception_cases = []
    for case in test_cases:
        result = await get_case_info(args.url, case)
        assert isinstance(result, dict)
        case_id = result['id']
        assert isinstance(case_id, int)
        if case_count:
            logger.info("Running test case %s (%d of %d)", case_id, case, case_count)
        else:
            logger.info("Debugging test case %s (%d)", case_id, case)
            logger.setLevel(logging.DEBUG)
        try:
            await run_case(args.url, case)
        except Exception:  # pylint: disable=broad-exception-caught
            logger.exception('  runtime exception during test case %s (%d)', case_id, case)
            exception_cases.append(case_id)
        logger.setLevel(logging.INFO)
    logger.info('Updating report')
    await update_reports(args.url)
    if exception_cases:
        logger.error('Runtime exception in %d of %d test cases: %s',
                     len(exception_cases), len(test_cases), exception_cases)
        sys.exit(1)


def parse_args() -> argparse.Namespace:
    ''' Parse command line arguments. '''
    parser = argparse.ArgumentParser(description='Autobahn client for'
        ' trio-websocket')
    parser.add_argument('url', help='WebSocket URL for server')
    # TODO: accept case ID's rather than indices
    parser.add_argument('debug_cases', type=int, nargs='*', help='Run'
        ' individual test cases with debug logging (optional)')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    trio.run(run_tests, args)
