#!/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 optparse
import subprocess
import sys
import threading
import time

stdout_lock = threading.Lock()
class FilterFormat:
  total_tests = 0
  finished_tests = 0

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

  def print_test_status(self, last_finished_test, time_ms):
    print "[%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]:
          print line
        print ("[%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])
      print "[0/%d] Running tests...\r" % 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:
      print "FAILED TESTS (%d/%d):" % (len(self.failures), self.total_tests)
      for (binary, test) in self.failures:
        print " ", 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

# 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=16,
                  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.
tests = []
for test_binary in binaries:
  command = [test_binary]
  if options.gtest_filter != '':
    command += ['--gtest_filter=' + options.gtest_filter]
  if options.gtest_also_run_disabled_tests:
    command += ['--gtest_also_run_disabled_tests']

  test_list = subprocess.Popen(command + ['--gtest_list_tests'],
                               stdout=subprocess.PIPE).communicate()[0]

  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

    test = test_group + line
    tests.append((test_binary, command, test))

# 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
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(str(job_id) + ': EXIT ' + str(code) + ' ' + str(runtime_ms))
  if code != 0:
    global exit_code
    exit_code = code

def worker():
  global job_id
  while True:
    job = None
    test_lock.acquire()
    if job_id < len(tests):
      (test_binary, command, test) = 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
    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()
sys.exit(exit_code)