#!/usr/bin/env python3

"""Run a command n times, and give stats on how many times it succeeds and fails."""

import itertools
import os
import shlex
import statistics
import subprocess
import sys
import time
import typing


def usage(retval: int):
    """Output a usage message."""
    if retval == 0:
        file_ = sys.stdout
    else:
        file_ = sys.stderr

    print('Run cmd repcount times, and gather simple stats about successes and failures', file=file_)
    print(file=file_)
    print(f'Usage: {sys.argv[0]}', file=file_)
    print('    --ignore-tty', file=file_)
    print('    --command cmd', file=file_)
    print('    --reps repcount', file=file_)
    print('    --verbose', file=file_)
    print('    --after-each', file=file_)
    print('    --example', file=file_)
    print('    --progress', file=file_)
    print('    --parameterize', file=file_)
    print('    --until-true', file=file_)
    print('    --until-false', file=file_)
    print('    --use-subshell', file=file_)
    print('    --help', file=file_)
    print(file=file_)
    print('--reps repcount says to run the command repcount times.  Conflicts with --until-true and --until-false', file=file_)
    print('--until-true says to run until the first "true" exit code is detected', file=file_)
    print('--until-false says to run until the first "false" exit code is detected', file=file_)
    print('--verbose says to dump output from all commands, successful or not', file=file_)
    print('--after-each says to write output from each command to stdout, without waiting until the end of the run', file=file_)
    print('--example says to output the last failed command run', file=file_)
    print('--examples says to output the last failed and last successful command run', file=file_)
    print('--progress says to count down to 0 as the reps are completed', file=file_)
    print('--parameterize says to replace {} with the repetition number', file=file_)
    print('--use-subshell says to spawn a subshell on "command" instead of using shlex.split()', file=file_)

    sys.exit(retval)


class Result:
    """A container class to hold results of a single run."""

    def __init__(self, command: typing.List[str], output: bytes, returncode: int, duration: float) -> None:
        """Initialize."""
        self.command = command
        # Saving the output could become expensive sometimes.
        self.output = output
        self.returncode = returncode
        self.exit_true = not bool(returncode)
        self.duration = duration

    def __str__(self):
        """Return a string version of this Result."""
        return '{} {}'.format(self.returncode, self.command)

    __repr__ = __str__


def output_one(result: Result, first_line: bytes):
    """Output one result."""
    os.write(1, first_line)
    os.write(1, result.output)
    os.write(1, b'Duration in seconds for this rep: %.1f\n' % (result.duration, ))


def output_one_fail(resultno: int, result: Result):
    """Output one failed run."""
    output_one(result=result, first_line=b'Failed run (run #%d):\n' % (resultno, ))


def output_one_success(resultno: int, result: Result):
    """Output one successful run."""
    output_one(result=result, first_line=b'Successful run (run #%d):\n' % (resultno, ))


def time_and_units(t: float) -> typing.Tuple[float, str]:
    """Convert a time in seconds to (possibly) minutes, or leave as seconds, depending on magnitude."""
    if t >= 60:
        return t / 60.0, 'minutes'
    else:
        return t, 'seconds'


def main() -> None:
    """Parse command line arguments, run the command n times, and summarize."""
    progress = False
    verbose = False
    after_each = False
    example = False
    examples = False
    reps = None
    command = None
    parameterize = False
    use_subshell = False
    until_true = False
    until_false = False
    warn_on_tty = True

    while sys.argv[1:]:
        if sys.argv[1] == '--command':
            command = sys.argv[2]
            del sys.argv[1]
        elif sys.argv[1] == '--reps':
            reps = int(sys.argv[2])
            del sys.argv[1]
        elif sys.argv[1] == '--until-true':
            until_true = True
        elif sys.argv[1] == '--until-false':
            until_false = True
        elif sys.argv[1] == '--example':
            example = True
        elif sys.argv[1] == '--examples':
            examples = True
        elif sys.argv[1] == '--parameterize':
            parameterize = True
        elif sys.argv[1] == '--use-subshell':
            use_subshell = True
        elif sys.argv[1] in ('-p', '--progress'):
            progress = True
        elif sys.argv[1] in ('-v', '--verbose'):
            verbose = True
        elif sys.argv[1] in '--after-each':
            after_each = True
        elif sys.argv[1] in ('-i', '--ignore-tty'):
            warn_on_tty = False
        elif sys.argv[1] in ('-h', '--help'):
            usage(0)
        else:
            print('{}: unrecognized option: {}'.format(sys.argv[0], sys.argv[1]), file=sys.stderr)
            usage(1)
        del sys.argv[1]

    if verbose + example + examples + after_each > 1:
        print('--verbose, --example, --examples and --after-each conflict', file=sys.stderr)
        usage(1)

    if command is None:
        print('--command is a required option', file=sys.stderr)
        usage(1)
    assert command is not None

    if (reps is not None) + until_true + until_false != 1:
        print('You must specify exactly one of --reps, --until-true and --until-false', file=sys.stderr)
        usage(1)

    if reps is None:
        # --reps not specified
        generator: typing.Union[range, typing.Iterator[int]] = itertools.count(0)
    else:
        generator = range(reps)

    if warn_on_tty and sys.stdout.isatty():
        print('Warning: stdout is a tty')

    results = []

    true_count = 0
    false_count = 0

    first_time = True

    time_so_far = 0.0

    very_beginning = time.time()

    try:
        for repno in generator:
            if first_time and ';' in command and not use_subshell:
                print('Warning: Found ; in command, but --use-subshell not specified - result may be puzzling')
                first_time = False

            if parameterize:
                cmd = command.replace('{}', str(repno))
            else:
                cmd = command

            if use_subshell:
                cmd2 = [cmd]
            else:
                cmd2 = shlex.split(cmd)

            start_of_rep = time.time()
            run_result = subprocess.run(
                cmd2,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
            )
            end_of_rep = time.time()

            time_for_this_rep = end_of_rep - start_of_rep

            if after_each:
                os.write(1, b'\n')
                os.write(1, run_result.stdout)

            result = Result(cmd2, run_result.stdout, run_result.returncode, duration=time_for_this_rep)

            time_so_far = time.time() - very_beginning

            results.append(result)

            if progress:
                if result.exit_true:
                    true_count += 1
                    if until_true:
                        break
                else:
                    false_count += 1
                    if until_false:
                        break
                if reps is not None:
                    rem_reps = reps - repno - 1
                curtime = time.ctime()
                avgtime = time_so_far / (repno + 1)
                # time_so_far / (repno + 1) == total_time / reps
                # estimated_secs_to_complete = time_so_far * reps / (repno + 1)
                if reps is not None:
                    eta = time.ctime(very_beginning + avgtime * reps)
                sys.stderr.write('\r')
                list_ = []
                if reps is not None:
                    list_.append(f'remaining: {rem_reps}')
                else:
                    list_.append(f'executed: {repno + 1}')
                list_.append(f'trues: {true_count}')
                list_.append(f'falses: {false_count}')
                list_.append(f'repstart: {curtime}')
                tm, units = time_and_units(avgtime)
                list_.append(f'avgtime: {tm:.1f} {units}')
                if reps is not None:
                    list_.append(f'eta: {eta}')
                files = [sys.stderr]
                if after_each:
                    files.append(sys.stdout)
                for file_ in files:
                    file_.write(', '.join(list_))
                    # with wildly varying avgtime, 7 may not always be enough - but usually it should be fine.
                    file_.write(' ' * 7)
                    if repno and (repno + 1) % 100 == 0 and repno + 1 != reps:
                        file_.write('\n')
                    file_.flush()
    except KeyboardInterrupt:
        got_ctrl_c = True
    else:
        got_ctrl_c = False

    if progress:
        sys.stderr.write('\n')

    if results:
        if examples:
            for resultno, result in reversed(list(enumerate(results))):
                if not result.exit_true:
                    continue
                output_one_success(resultno, result)
                break
            else:
                print('No successful examples found')

            for resultno, result in reversed(list(enumerate(results))):
                if result.exit_true:
                    continue
                output_one_fail(resultno, result)
                break
            else:
                print('No fail examples found')

        if example:
            for resultno, result in reversed(list(enumerate(results))):
                if result.exit_true:
                    continue
                output_one_fail(resultno, result)
                break
            else:
                print('No fail examples found')

        if verbose:
            for resultno, result in enumerate(results):
                if result.exit_true:
                    output_one_success(resultno, result)
                else:
                    output_one_fail(resultno, result)

        mean_truth = statistics.mean(result.exit_true for result in results)
        mean_duration = statistics.mean(result.duration for result in results)
        print(f'That was {len(results)} repetitions', file=sys.stderr)
        print(f'Exited true {mean_truth * 100.0:.1f}% of the time', file=sys.stderr)
        tm, units = time_and_units(mean_duration)
        print(f'Mean duration: {tm:.1f} {units}', file=sys.stderr)

    if got_ctrl_c:
        raise SystemExit('Got control-c')

    if all(result.exit_true for result in results):
        sys.exit(0)
    else:
        sys.exit(1)


if __name__ == '__main__':
    main()