import argparse
import os
import pathlib
import signal
import subprocess
import sys

import pytest


SCRIPT_DIR = pathlib.Path(__file__).parent.resolve()
SUITE_TIMEOUT_S = 120


class PassedCounter:
    def __init__(self):
        self.passed = 0

    def pytest_report_teststatus(self, report, config):
        if report.when == 'call' and report.passed:
            self.passed += 1


def suite_score(passed, expected, maximum, exit_code=pytest.ExitCode.OK):
    return maximum if exit_code == pytest.ExitCode.OK and passed == expected else 0


def run_suite(filename, expected, maximum):
    # Each component has its own deadline so a timeout cannot discard other scores.
    process = subprocess.Popen(
        [sys.executable, '-u', str(pathlib.Path(__file__).resolve()),
         '--suite', str(SCRIPT_DIR / filename), str(expected)],
        start_new_session=os.name == 'posix',
    )
    try:
        exit_code = process.wait(timeout=SUITE_TIMEOUT_S)
    except subprocess.TimeoutExpired:
        if os.name == 'posix':
            # Include grpcurl and other subprocesses started by this test suite.
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass  # The suite may have exited just after wait() timed out.
        else:
            process.kill()
        process.wait()
        print(f'{filename} exceeded {SUITE_TIMEOUT_S} seconds; component score is zero.', flush=True)
        return 0
    return maximum if exit_code == 0 else 0


def suite_exit_code(filename, expected):
    counter = PassedCounter()
    exit_code = pytest.main(['-vs', filename], plugins=[counter])
    return 0 if suite_score(counter.passed, expected, 1, exit_code) else 1


def component_enabled(environment_name):
    return os.environ.get(environment_name, '1') == '1'


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--component',
        choices=('all', 'proto', 'server', 'client'),
        default='all',
    )
    parser.add_argument('--suite', nargs=2, metavar=('FILE', 'EXPECTED'), help=argparse.SUPPRESS)
    return parser.parse_args()


def main():
    args = parse_args()
    if args.suite is not None:
        return suite_exit_code(args.suite[0], int(args.suite[1]))
    component = args.component
    score = 0

    if component in ('all', 'proto'):
        proto_score = run_suite('test_proto.py', expected=1, maximum=2)
        score += proto_score
        print(f'Proto: {proto_score}/2')
        print()

    if component in ('all', 'server'):
        if component_enabled('MESSENGER_SERVER_TESTS_ENABLED'):
            server_score = run_suite('test_server.py', expected=4, maximum=4)
        else:
            server_score = 0
            print('Server tests were not run because the server image did not build.')
        score += server_score
        print(f'Server: {server_score}/4')
        print()

    if component in ('all', 'client'):
        if component_enabled('MESSENGER_CLIENT_TESTS_ENABLED'):
            client_score = run_suite('test_client.py', expected=3, maximum=4)
        else:
            client_score = 0
            print('Client tests were not run because the client image did not build.')
        score += client_score
        print(f'Client: {client_score}/4')

    print(f'\nSCORE: {score}')


if __name__ == '__main__':
    sys.exit(main())