#!/usr/bin/env python3
import contextlib
import argparse
import errno
import io
import json
from pathlib import Path
import subprocess
import tempfile
import time

import agent_session


def check(binary, options, command, expected, phase=None, cleanup_failure=False):
    children = []
    original_popen = subprocess.Popen

    def launch(*args, **kwargs):
        child = original_popen(*args, **kwargs)
        children.append(child)
        return child

    with tempfile.TemporaryDirectory(prefix='pardes-agent-driver-test-') as directory:
        original_cleanup = tempfile.TemporaryDirectory.cleanup

        def fail_cleanup(temporary):
            original_cleanup(temporary)
            raise OSError(errno.ENOTEMPTY, 'injected cleanup failure')

        previous_tempdir = tempfile.tempdir
        tempfile.tempdir = directory
        subprocess.Popen = launch
        if cleanup_failure:
            tempfile.TemporaryDirectory.cleanup = fail_cleanup
        stdout, stderr = io.StringIO(), io.StringIO()
        started = time.monotonic()
        try:
            with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
                try:
                    result = agent_session.main([binary, *options, '--', *command])
                except SystemExit as error:
                    result = error.code
        finally:
            subprocess.Popen = original_popen
            tempfile.tempdir = previous_tempdir
            tempfile.TemporaryDirectory.cleanup = original_cleanup
        assert result == expected, (result, expected, stderr.getvalue())
        assert all(child.poll() is not None for child in children), 'owned Pardes still running'
        assert not list(Path(directory).iterdir()), 'owned session files remain'
        if expected == 2:
            assert not children, 'invalid arguments started Pardes'
        elif expected == 0:
            metrics = json.loads(stdout.getvalue())
            assert metrics['body_bytes'] > 0 and metrics['body_rows'] > 0
            assert len(metrics['visible_text_sha256']) == len(metrics['body_sha256']) == 64
            assert metrics['ninep_write_to_observation_ms'] >= 0
            assert metrics['ninep_clear_to_observation_ms'] >= 0
            assert 'AGENT_READY' not in stdout.getvalue()
        else:
            assert not stdout.getvalue()
            error = json.loads(stderr.getvalue())
            assert error['phase'] == phase
            if cleanup_failure:
                assert error['error'] == 'OSError' and error['errno'] == errno.ENOTEMPTY
                assert error['metrics']['body_rows'] > 0
                assert error['expected_visible_text_sha256'] == error['observed_visible_text_sha256']
            else:
                assert error['error'] == 'TimeoutError'
            assert 'AGENT_READY' not in stderr.getvalue()
            assert time.monotonic() - started < 8, 'timeout did not bound cleanup'


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('binary', type=Path)
    parser.add_argument('--gui-grid', action='store_true')
    args = parser.parse_args()
    binary = str(args.binary.resolve())
    ready = ['--ready', 'AGENT_READY', *(['--gui-grid'] if args.gui_grid else [])]
    command = ['/bin/sh', '-c', 'printf "AGENT_READY> "; read answer']
    for options, argv in [
        ([], command),
        (ready + ['--timeout', 'nan'], command),
        (ready + ['--timeout', '0'], command),
        (ready + ['--timeout', '601'], command),
        (ready + ['--min-rows', '-1'], command),
        (ready + ['--prepare-key', '0a'], command),
        (ready + ['--prepare-key', '0d'], command),
        (ready + ['--prepare-key', '41'], command),
        (ready + ['--prepare-key', '0f0f'], command),
        (ready, []),
    ]:
        check(binary, options, argv, 2)
    check(binary, ready + ['--timeout', '10'], command, 0)
    check(binary, ready + ['--timeout', '10', '--min-rows', '1000'],
          ['/bin/sh', '-c', 'printf "AGENT_READY> "; sleep .3; seq 1 3000; printf "AGENT_READY> "; read answer'], 0)
    check(binary, ready + ['--timeout', '10', '--min-rows', '1000', '--prepare-key', '0f'],
          ['/bin/sh', '-c', 'stty raw -echo; printf "AGENT_READY> "; dd bs=1 count=1 of=/dev/null 2>/dev/null; '
           'stty sane; seq 1 3000; printf "AGENT_READY> "; read answer'], 0)
    check(binary, ready + ['--timeout', '2', '--min-rows', '1000'], command, 1, 'history')
    check(binary, ready + ['--timeout', '1'],
          ['/bin/sh', '-c', ': AGENT_READY; sleep 30'], 1, 'readiness')
    check(binary, ready + ['--timeout', '10'], command, 1, 'files-cleanup', cleanup_failure=True)
    print('agent-session: 16 functional argument, interaction, timeout and cleanup checks passed')


if __name__ == '__main__':
    main()