#!/usr/bin/env python2
# Copyright 2013 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import cPickle
import gzip
import multiprocessing
import optparse
import os
import subprocess
import sys
import threading
import time
import zlib

# Return the width of the terminal, or None if it couldn't be
# determined (e.g. because we're not being run interactively).
def term_width(out):
  if not out.isatty():
    return None
  try:
    p = subprocess.Popen(["stty", "size"],
                         stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = p.communicate()
    if p.returncode != 0 or err:
      return None
    return int(out.split()[1])
  except (IndexError, OSError, ValueError):
    return None

# Output transient and permanent lines of text. If several transient
# lines are written in sequence, the new will overwrite the old. We
# use this to ensure that lots of unimportant info (tests passing)
# won't drown out important info (tests failing).
class Outputter(object):
  def __init__(self, out_file):
    self.__out_file = out_file
    self.__previous_line_was_transient = False
    self.__width = term_width(out_file)  # Line width, or None if not a tty.
  def transient_line(self, msg):
    if self.__width is None:
      self.__out_file.write(msg + "\n")
    else:
      self.__out_file.write("\r" + msg[:self.__width].ljust(self.__width))
      self.__previous_line_was_transient = True
  def permanent_line(self, msg):
    if self.__previous_line_was_transient:
      self.__out_file.write("\n")
      self.__previous_line_was_transient = False
    self.__out_file.write(msg + "\n")

stdout_lock = threading.Lock()

class FilterFormat:
  out = Outputter(sys.stdout)
  total_tests = 0
  finished_tests = 0

  tests = {}
  outputs = {}
  failures = []

  def print_test_status(self, last_finished_test, time_ms):
    self.out.transient_line("[%d/%d] %s (%d ms)"
                            % (self.finished_tests, self.total_tests,
                               last_finished_test, time_ms))

  def handle_meta(self, job_id, args):
    (command, arg) = args.split(' ', 1)
    if command == "TEST":
      (binary, test) = arg.split(' ', 1)
      self.tests[job_id] = (binary, test.strip())
      self.outputs[job_id] = []
    elif command == "EXIT":
      (exit_code, time_ms) = [int(x) for x in arg.split(' ', 1)]
      self.finished_tests += 1
      (binary, test) = self.tests[job_id]
      self.print_test_status(test, time_ms)
      if exit_code != 0:
        self.failures.append(self.tests[job_id])
        for line in self.outputs[job_id]:
          self.out.permanent_line(line)
        self.out.permanent_line(
          "[%d/%d] %s returned/aborted with exit code %d (%d ms)"
          % (self.finished_tests, self.total_tests, test, exit_code, time_ms))
    elif command == "TESTCNT":
      self.total_tests = int(arg.split(' ', 1)[1])
      self.out.transient_line("[0/%d] Running tests..." % self.total_tests)

  def add_stdout(self, job_id, output):
    self.outputs[job_id].append(output)

  def log(self, line):
    stdout_lock.acquire()
    (prefix, output) = line.split(' ', 1)

    if prefix[-1] == ':':
      self.handle_meta(int(prefix[:-1]), output)
    else:
      self.add_stdout(int(prefix[:-1]), output)
    stdout_lock.release()

  def end(self):
    if self.failures:
      self.out.permanent_line("FAILED TESTS (%d/%d):"
                              % (len(self.failures), self.total_tests))
      for (binary, test) in self.failures:
        self.out.permanent_line(" " + binary + ": " + test)

class RawFormat:
  def log(self, line):
    stdout_lock.acquire()
    sys.stdout.write(line + "\n")
    sys.stdout.flush()
    stdout_lock.release()
  def end(self):
    pass

# Record of test runtimes. Has built-in locking.
class TestTimes(object):
  def __init__(self, save_file):
    "Create new object seeded with saved test times from the given file."
    self.__times = {}  # (test binary, test name) -> runtime in ms

    # Protects calls to record_test_time(); other calls are not
    # expected to be made concurrently.
    self.__lock = threading.Lock()

    try:
      with gzip.GzipFile(save_file, "rb") as f:
        times = cPickle.load(f)
    except (EOFError, IOError, cPickle.UnpicklingError, zlib.error):
      # File doesn't exist, isn't readable, is malformed---whatever.
      # Just ignore it.
      return

    # Discard saved times if the format isn't right.
    if type(times) is not dict:
      return
    for ((test_binary, test_name), runtime) in times.items():
      if (type(test_binary) is not str or type(test_name) is not str
          or type(runtime) not in {int, long}):
        return

    self.__times = times

  def get_test_time(self, binary, testname):
    "Return the last duration for the given test, or 0 if there's no record."
    return self.__times.get((binary, testname), 0)

  def record_test_time(self, binary, testname, runtime_ms):
    "Record that the given test ran in the specified number of milliseconds."
    with self.__lock:
      self.__times[(binary, testname)] = runtime_ms

  def write_to_file(self, save_file):
    "Write all the times to file."
    try:
      with open(save_file, "wb") as f:
        with gzip.GzipFile("", "wb", 9, f) as gzf:
          cPickle.dump(self.__times, gzf, cPickle.HIGHEST_PROTOCOL)
    except IOError:
      pass  # ignore errors---saving the times isn't that important

# Remove additional arguments (anything after --).
additional_args = []

for i in range(len(sys.argv)):
  if sys.argv[i] == '--':
    additional_args = sys.argv[i+1:]
    sys.argv = sys.argv[:i]
    break

parser = optparse.OptionParser(
    usage = 'usage: %prog [options] binary [binary ...] -- [additional args]')

parser.add_option('-r', '--repeat', type='int', default=1,
                  help='repeat tests')
parser.add_option('-w', '--workers', type='int',
                  default=multiprocessing.cpu_count(),
                  help='number of workers to spawn')
parser.add_option('--gtest_color', type='string', default='yes',
                  help='color output')
parser.add_option('--gtest_filter', type='string', default='',
                  help='test filter')
parser.add_option('--gtest_also_run_disabled_tests', action='store_true',
                  default=False, help='run disabled tests too')
parser.add_option('--format', type='string', default='filter',
                  help='output format (raw,filter)')

(options, binaries) = parser.parse_args()

if binaries == []:
  parser.print_usage()
  sys.exit(1)

logger = RawFormat()
if options.format == 'raw':
  pass
elif options.format == 'filter':
  logger = FilterFormat()
else:
  sys.exit("Unknown output format: " + options.format)

# Find tests.
save_file = os.path.join(os.path.expanduser("~"), ".gtest-parallel-times")
times = TestTimes(save_file)
tests = []
for test_binary in binaries:
  command = [test_binary]
  if options.gtest_also_run_disabled_tests:
    command += ['--gtest_also_run_disabled_tests']

  list_command = list(command)
  if options.gtest_filter != '':
    list_command += ['--gtest_filter=' + options.gtest_filter]

  try:
    test_list = subprocess.Popen(list_command + ['--gtest_list_tests'],
                                 stdout=subprocess.PIPE).communicate()[0]
  except OSError as e:
    sys.exit("%s: %s" % (test_binary, str(e)))

  command += additional_args

  test_group = ''
  for line in test_list.split('\n'):
    if not line.strip():
      continue
    if line[0] != " ":
      test_group = line.strip()
      continue
    line = line.strip()
    if not options.gtest_also_run_disabled_tests and 'DISABLED' in line:
      continue
    line = line.split('#')[0].strip()
    if not line:
      continue

    test = test_group + line
    tests.append((times.get_test_time(test_binary, test),
                  test_binary, test, command))
tests.sort(reverse=True)

# Repeat tests (-r flag).
tests *= options.repeat
test_lock = threading.Lock()
job_id = 0
logger.log(str(-1) + ': TESTCNT ' + ' ' + str(len(tests)))

exit_code = 0

# Run the specified job. Return the elapsed time in milliseconds if
# the job succeeds, or a very large number (larger than any reasonable
# elapsed time) if the job fails. (This ensures that failing tests
# will run first the next time.)
def run_job((command, job_id, test)):
  begin = time.time()
  sub = subprocess.Popen(command + ['--gtest_filter=' + test] +
                           ['--gtest_color=' + options.gtest_color],
                         stdout = subprocess.PIPE,
                         stderr = subprocess.STDOUT)

  while True:
    line = sub.stdout.readline()
    if line == '':
      break
    logger.log(str(job_id) + '> ' + line.rstrip())

  code = sub.wait()
  runtime_ms = int(1000 * (time.time() - begin))
  logger.log("%s: EXIT %s %d" % (job_id, code, runtime_ms))
  if code == 0:
    return runtime_ms
  global exit_code
  exit_code = code
  return sys.maxint

def worker():
  global job_id
  while True:
    job = None
    test_lock.acquire()
    if job_id < len(tests):
      (_, test_binary, test, command) = tests[job_id]
      logger.log(str(job_id) + ': TEST ' + test_binary + ' ' + test)
      job = (command, job_id, test)
    job_id += 1
    test_lock.release()
    if job is None:
      return
    times.record_test_time(test_binary, test, run_job(job))

def start_daemon(func):
  t = threading.Thread(target=func)
  t.daemon = True
  t.start()
  return t

workers = [start_daemon(worker) for i in range(options.workers)]

[t.join() for t in workers]
logger.end()
times.write_to_file(save_file)
sys.exit(exit_code)