#!/usr/bin/env python3

"""Dispatch to update_*_test_checks.py scripts automatically in bulk

Given a list of test files, this script will invoke the correct
update_test_checks-style script, skipping any tests which have not previously
had assertions autogenerated.
"""

from __future__ import print_function

import argparse
import os
import re
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor

RE_ASSERTIONS = re.compile(
    r'NOTE: Assertions have been autogenerated by ([^\s]+)( UTC_ARGS:.*)?$')

def find_utc_tool(search_path, utc_name):
  """
  Return the path to the given UTC tool in the search path, or None if not
  found.
  """
  for path in search_path:
    candidate = os.path.join(path, utc_name)
    if os.path.isfile(candidate):
      return candidate
  return None

def run_utc_tool(utc_name, utc_tool, testname):
  result = subprocess.run([utc_tool, testname], stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE)
  return (result.returncode, result.stdout, result.stderr)

def main():
  from argparse import RawTextHelpFormatter
  parser = argparse.ArgumentParser(description=__doc__,
                                   formatter_class=RawTextHelpFormatter)
  parser.add_argument(
      '--jobs', '-j', default=1, type=int,
      help='Run the given number of jobs in parallel')
  parser.add_argument(
      '--utc-dir', nargs='*',
      help='Additional directories to scan for update_*_test_checks scripts')
  parser.add_argument('tests', nargs='+')
  config = parser.parse_args()

  if config.utc_dir:
    utc_search_path = config.utc_dir[:]
  else:
    utc_search_path = []
  script_name = os.path.abspath(__file__)
  utc_search_path.append(os.path.join(os.path.dirname(script_name),
                                      os.path.pardir))

  not_autogenerated = []
  utc_tools = {}
  have_error = False

  with ThreadPoolExecutor(max_workers=config.jobs) as executor:
    jobs = []

    for testname in config.tests:
      with open(testname, 'r') as f:
        header = f.readline().strip()
        m = RE_ASSERTIONS.search(header)
        if m is None:
          not_autogenerated.append(testname)
          continue

        utc_name = m.group(1)
        if utc_name not in utc_tools:
          utc_tools[utc_name] = find_utc_tool(utc_search_path, utc_name)
          if not utc_tools[utc_name]:
            print(f"{utc_name}: not found (used in {testname})",
                  file=sys.stderr)
            have_error = True
            continue

        future = executor.submit(run_utc_tool, utc_name, utc_tools[utc_name],
                                 testname)
        jobs.append((testname, future))

    for testname, future in jobs:
      return_code, stdout, stderr = future.result()

      print(f"Update {testname}")
      stdout = stdout.decode(errors='replace')
      if stdout:
        print(stdout, end='')
        if not stdout.endswith('\n'):
          print()

      stderr = stderr.decode(errors='replace')
      if stderr:
        print(stderr, end='')
        if not stderr.endswith('\n'):
          print()
      if return_code != 0:
        print(f"Return code: {return_code}")
        have_error = True

  if have_error:
    sys.exit(1)

  if not_autogenerated:
    print("Tests without autogenerated assertions:")
    for testname in not_autogenerated:
      print(f"  {testname}")

if __name__ == '__main__':
  main()