import argparse
import os
import subprocess
import sys

PYTHON_CMD = 'python'

BROWSER_TEST_CMD = os.path.join('testing', 'scripts',
    'run_gpu_integration_test_as_googletest.py')
GPU_TEST_CMD     = os.path.join('content', 'test', 'gpu',
    'run_gpu_integration_test.py')
ANGLE_TEST_CMD   = 'angle_end2end_tests'

WEBGL_TEST_OPT = 'webgl_conformance'
TRACE_TEST_OPT = 'trace_test'
INFO_TEST_OPT  = 'info_collection'

WEBGL2_CONFORMANCE_VERSION = '2.0.1'
WEBGL2_ABBREVIATED_RESULT  = os.path.join('content', 'test', 'data', 'gpu',
    'webgl2_conformance_tests_output.json')

TRYBOT_SHARDS = {
  'webgl' : 2,
  'webgl2' : 20,
  'angle' : 4,
  'info' : 1,
  'trace' : 1
}

RESERVED_OPTIONS = [
  '-v', '--verbose',
  '--show-stdout', '--passthrough',
  '--browser', '--browser-executable=',
  '--chrome-root', '--chromium-output-directory', '--top-level-dirs',
  '--extra-browser-args', '--enable-logging', '--js-flags',
  '--use-angle', '--use-gl', '--use-cmd-decoder',
  '--write-full-results-to', '--isolated-script-test-output',
  '--webgl-conformance-version', '--read-abbreviated-json-results-from',
  '--expected-vendor-id', '--expected-device-id',
  '--total-shards', '--shard-index',
  '--test-launcher-total-shards', '--test-launcher-shard-index',
  '--test-filter', '--gtest_filter',
  '--test-launcher-retry-limit',
  '--use-gpu-in-tests'
]

def execute_cmd(cmd, log=None):
  proc = subprocess.Popen(cmd,
      stdout=subprocess.PIPE if log else None, stderr=subprocess.STDOUT,
      shell=sys.platform.startswith('win'))

  if log:
    with open(log, 'w') as log_file:
      for line in iter(proc.stdout.readline, b''):
        line = line.strip()
        print(line)
        log_file.write(line + '\n')

  return proc.wait()

def execute_shard(cmd, args):
  shard_postfix = ''
  if args.shard > 1:
    shard_postfix = '.shard' + format(args.index, '02d')

  result_name, result_ext = os.path.splitext(args.result)
  log_name, log_ext = os.path.splitext(args.log)

  for n in range(1, args.repeat+1):
    postfix = shard_postfix
    if args.repeat > 1:
      postfix += '.' + format(n, '03d')

    new_cmd = cmd[:]
    if not args.type == 'angle':
      new_cmd.append('--isolated-script-test-output='
          + result_name + postfix + result_ext)

    print('\n' + ' '.join(new_cmd) + '\n')
    if args.print_only:
      return

    log = log_name + postfix + log_ext
    execute_cmd(new_cmd, log)

def update_global_arguments(args):
  global BROWSER_TEST_CMD
  global GPU_TEST_CMD
  global ANGLE_TEST_CMD
  global WEBGL2_ABBREVIATED_RESULT

  if args.build == 'default':
    ANGLE_TEST_CMD = os.path.join('out', 'Default', ANGLE_TEST_CMD)
  elif args.build == 'release':
    ANGLE_TEST_CMD = os.path.join('out', 'Release', ANGLE_TEST_CMD)
  elif args.build == 'debug':
    ANGLE_TEST_CMD = os.path.join('out', 'Debug', ANGLE_TEST_CMD)

  if args.dir:
    BROWSER_TEST_CMD = os.path.join(args.dir, BROWSER_TEST_CMD)
    GPU_TEST_CMD = os.path.join(args.dir, GPU_TEST_CMD)
    ANGLE_TEST_CMD = os.path.join(args.dir, ANGLE_TEST_CMD)
    WEBGL2_ABBREVIATED_RESULT = os.path.join(args.dir, WEBGL2_ABBREVIATED_RESULT)

def generate_webgl_arguments(args):
  # Common arguments
  common_args = ['--show-stdout', '--browser=' + args.build,
      '--passthrough', '-v']
  if args.dir :
    common_args.append('--chrome-root=' + args.dir)

  # Browser arguments
  # TODO: check --enable-logging=stderr
  browser_args = ['--js-flags=--expose-gc', '--enable-logging=stderr']
  if args.backend == 'd3d9v':
    browser_args.extend(['--use-angle=d3d9', '--use-gl=angle',
        '--use-cmd-decoder=validating'])
  elif args.backend == 'd3d11v':
    browser_args.extend(['--use-angle=d3d11',
        '--use-cmd-decoder=validating'])
  elif args.backend == 'd3d9':
    browser_args.extend(['--use-angle=d3d9',
        '--use-cmd-decoder=passthrough'])
  elif args.backend == 'd3d11':
    browser_args.extend(['--use-angle=d3d11',
        '--use-cmd-decoder=passthrough'])
  elif args.backend == 'gl':
    browser_args.extend(['--use-angle=gl', '--use-gl=angle',
        '--use-cmd-decoder=passthrough'])
  elif args.backend == 'vulkan':
    browser_args.extend(['--use-angle=vulkan',
        '--use-cmd-decoder=passthrough'])

  # WebGL arguments
  webgl_args = []
  if args.type == 'webgl2':
    webgl_args.append('--webgl-conformance-version='
        + WEBGL2_CONFORMANCE_VERSION)
    webgl_args.append('--read-abbreviated-json-results-from='
        + WEBGL2_ABBREVIATED_RESULT)
  elif args.type == 'info':
    webgl_args.extend(['--expected-vendor-id', '8086'])
    webgl_args.extend(['--expected-device-id', '5912'])

  total_args = []
  total_args.extend(common_args)
  total_args.append('--extra-browser-args=' + ' '.join(browser_args))
  total_args.extend(webgl_args)

  if args.filter:
      filter = args.filter
      if (not filter.startswith('*') and not filter.startswith('deqp')
          and not filter.startswith('conformance')):
        filter = '*' + filter
      if not filter.endswith('*') and not filter.endswith('html'):
        filter = filter + '*'
      total_args.append('--test-filter=' + filter)

  return total_args

def generate_angle_arguments(args):
  total_args = ['--use-gpu-in-tests', '--test-launcher-retry-limit=0']
  if args.filter:
    total_args.append('--gtest_filter=' + args.filter)
  return total_args

def main():
  parser = argparse.ArgumentParser(
      description='Run WebGL and ANGLE tests',
      formatter_class=argparse.RawTextHelpFormatter)
  parser.add_argument('type',
      choices=['webgl', 'webgl2', 'angle', 'trace', 'info'],
      help='Specify the test you want to run.\n\n'\
           'webgl   webgl conformance tests\n'\
           'webgl2  webgl2 conformance tests\n'\
           'angle   angle end2end tests\n'\
           'trace   trace tests\n'\
           'info    info collection tests')
  parser.add_argument('--backend', '-e',
      choices=['d3d9v', 'd3d11v', 'd3d9', 'd3d11', 'gl', 'vulkan'], default='',
      help='Specify the backend. Only available for webgl/webgl2.\n'\
           'Default is empty, will run default conformance tests.\n\n'\
           'd3d9v   d3d9 validating\n'\
           'd3d11v  d3d11 validating\n'\
           'd3d9    d3d9 passthrough\n'\
           'd3d11   d3d11 passthrough\n'\
           'gl      gl passthrough\n'\
           'vulkan  vulkan passthrough\n\n')
  parser.add_argument('--dir', '-d', default='',
      help='Chromium source directory.\n'\
           'Default is current directory.\n\n')
  parser.add_argument('--build', '-b',
      choices=['release', 'debug', 'default'], default='release',
      help='Browser type to run. Default is release.\n'\
           'release/debug/default assume that chrome executables are\n'\
           'generated into out/Release or out/Debug or out/Default.\n\n')
  parser.add_argument('--repeat', '-r', default='1', type=int,
      help='The number of times to repeat running this test.\n'\
           'If the number of shards is more than 1, the running sequence will\n'\
           'be shard0 * N times, shard1 * N times ...\n\n')
  parser.add_argument('--shard', '-s', default='1', type=int,
      help='Total number of shards being used for this test.\n'\
           'Default is 1.\n\n')
  parser.add_argument('--index', '-i', default='-1', type=int,
      help='Shard index of this test.\n'\
           'If the number of shards is more than 1 and this argument is not\n'\
           'specified, all shards will be ran in sequence.\n\n')
  parser.add_argument('--trybot', '-t', action='store_true',
      help='Use the same number of shards as official try bot.\n\n')
  parser.add_argument('--log', '-g', default='',
      help='Saves the test log to that path.\n'\
           'Default is [webgl2|d3d9|...]_test.log.\n\n')
  parser.add_argument('--result', '-w', default='',
      help='Writes the trybot style test results to a JSON file.\n'\
           'Default is [webgl2|d3d9|...]_test_result.json.\n\n')
  parser.add_argument('--filter', '-f', default='',
      help='Keyword to match the test cases.\n'\
           'Devide with |.\n\n')
  parser.add_argument('--print-only', '-p', action='store_true',
      help='Print the command only without executing it.\n\n')
  # parser.add_argument('--actual-result', '-a', default='',
  #     help='Writes the actual results and expected results to a JSON file.\n'\
  #          'Default is [webgl2|d3d9|...]_test__result_actual.json.\n\n')
  args, extra_args = parser.parse_known_args()

  # Check arguments
  if args.trybot:
    if args.shard > 1 or args.index > -1:
      sys.exit('Cannot specify shard and index when trybot is enabled')
    args.shard = TRYBOT_SHARDS[args.type]

  if args.shard < 1:
    sys.exit('Invalid shard number')
  elif args.index >= args.shard:
    sys.exit('Invalid index number')

  if args.repeat < 1:
    sys.exit('Invalid repeat number')

  if args.backend:
    if not args.type == 'webgl' and not args.type == 'webgl2':
      sys.exit('Backend is only available for webgl/webgl2')

  for arg in extra_args:
    for option in RESERVED_OPTIONS:
      if arg.startswith(option):
        sys.exit('This option has been reserved for internal usage: ' + arg)

  # Update arguments
  update_global_arguments(args)

  if not args.log:
    prefix = args.type
    if args.backend:
      prefix = prefix + '_' + args.backend
    args.log = prefix + '_test.log'

  if not args.result:
    prefix = args.type
    if args.backend:
      prefix = prefix + '_' + args.backend
    args.result = prefix + '_test_result.json'

  # Generate command
  cmd = []
  shard_arg = ''
  index_arg = ''

  if args.type == 'angle':
    cmd.append(ANGLE_TEST_CMD)
    cmd.extend(generate_angle_arguments(args))

    shard_arg = '--test-launcher-total-shards='
    index_arg = '--test-launcher-shard-index='
  else:
    cmd.append(PYTHON_CMD)
    cmd.append(BROWSER_TEST_CMD)
    cmd.append(GPU_TEST_CMD)
    if args.type == 'trace':
      cmd.append(TRACE_TEST_OPT)
    elif args.type == 'info':
      cmd.append(INFO_TEST_OPT)
    else:
      cmd.append(WEBGL_TEST_OPT)
    cmd.extend(generate_webgl_arguments(args))

    shard_arg = '--total-shards='
    index_arg = '--shard-index='

  cmd.extend(extra_args)

  # Execute command
  if args.shard == 1:
    execute_shard(cmd, args)
  else:
    cmd.append(shard_arg + str(args.shard))

    if args.index >= 0:
      cmd.append(index_arg + str(args.index))
      execute_shard(cmd, args)
    else:
      for i in range(0, args.shard):
        shard_cmd = cmd[:]
        shard_cmd.append(index_arg + str(i))
        args.index = i
        execute_shard(shard_cmd, args)

  return 0

if __name__ == '__main__':
  sys.exit(main())