# Copyright 2018 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Yet another domain specific client for Swarming."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import json
import os
import urllib

from lucifer import autotest
from skylab_staging import errors

# This is hard-coded everywhere -- on builders requesting skylab suite as well
# as in commands to be generated by users requesting one off suites.
_SKYLAB_RUN_SUITE_PATH = '/usr/local/autotest/bin/run_suite_skylab'
_SWARMING_POOL_SKYLAB_BOTS = 'ChromeOSSkylab'
_SWARMING_POOL_SKYLAB_SUITE_BOTS = 'ChromeOSSkylab-suite'
# Test push creates all suites at the highest allowed non-admin task priority.
# This ensures that test push tasks are prioritized over any user created tasks
# in the staging lab.
_TEST_PUSH_SUITE_PRIORITY = 50


class Client(object):
  """A domain specific client for Swarming service."""

  def __init__(self, cli_path, host, service_account_json=None):
    self._cli_path = cli_path
    self._host = host
    self._service_account_json = service_account_json

  def num_ready_duts(self, board, pool):
    """Count the number of DUTs in the given board, pool in dut_state ready.

    @param board: The board autotest label of the DUTs.
    @param pool: The pool autotest label of the DUTs.
    @returns number of DUTs in dut_state ready.
    """
    qargs = [
        ('dimensions', 'pool:%s' % _SWARMING_POOL_SKYLAB_BOTS),
        ('dimensions', 'label-board:%s' % board),
        ('dimensions', 'label-pool:%s' % pool),
        ('dimensions', 'dut_state:ready'),
    ]
    result = self.query('bots/count', qargs)
    if not result:
      return 0
    return int(result['count']) - (int(result['busy'])
                                   + int(result['dead'])
                                   + int(result['quarantined'])
                                   + int(result['maintenance']))

  def query(self, path, qargs):
    """Run a Swarming 'query' call.

    @param path: Path of the query RPC call.
    @qargs: Arguments for the RPC call.
    @returns: json response from the Swarming call.
    """
    cros_build_lib = autotest.chromite_load('cros_build_lib')
    cmdarg = path
    if qargs:
      cmdarg += "?%s" % urllib.urlencode(qargs)

    cmd = self._base_cmd('query') + [cmdarg]

    result = cros_build_lib.RunCommand(cmd, capture_output=True)
    return json.loads(result.output)

  def trigger_suite(self, board, pool, build, suite_name, timeout_s):
    """Trigger an autotest suite. Use wait_for_suite to wait for results.

    @param board: The board autotest label of the DUTs.
    @param pool: The pool autotest label of the DUTs.
    @param build: The build to test, e.g. link-paladin/R70-10915.0.0-rc1.
    @param suite_name: The name of the suite to run, e.g. provision.
    @param timeout_s: Timeout for the suite, in seconds.
    @returns: The task ID of the kicked off suite.
    """
    raw_cmd = self._suite_cmd_common(board, pool, build, suite_name, timeout_s)
    raw_cmd += ['--create_and_return']
    return self._run(board, pool, build, suite_name, timeout_s, raw_cmd)

  def wait_for_suite(self, task_id, board, pool, build, suite_name, timeout_s):
    """Wait for a suite previously kicked off via trigger_suite().

    @param task_id: Task ID of the suite, as returned by trigger_suite().
    @param board: The board autotest label of the DUTs.
    @param pool: The pool autotest label of the DUTs.
    @param build: The build to test, e.g. link-paladin/R70-10915.0.0-rc1.
    @param suite_name: The name of the suite to run, e.g. provision.
    @param timeout_s: Timeout for the suite, in seconds.
    @returns: The task ID of the kicked off suite.
    """
    raw_cmd = self._suite_cmd_common(board, pool, build, suite_name, timeout_s)
    raw_cmd += ['--suite_id', task_id]
    return self._run(board, pool, build, suite_name, timeout_s, raw_cmd)

  def _suite_cmd_common(self, board, pool, build, suite_name, timeout_s):
    return [
        _SKYLAB_RUN_SUITE_PATH,
        '--board', board,
        '--build', build,
        '--max_retries', '5',
        '--pool', _old_style_pool_label(pool),
        '--priority', str(_TEST_PUSH_SUITE_PRIORITY),
        '--suite_args', json.dumps({'num_required': 1}),
        '--suite_name', suite_name,
        '--test_retry',
        '--timeout_mins', str(int(timeout_s / 60)),
    ]

  def _run(self, board, pool, build, suite_name, timeout_s, raw_cmd):
    timeout_s = str(int(timeout_s))
    task_name = '%s-%s' % (build, suite_name)
    # This is a subset of the tags used by builders when creating suites.
    # These tags are used by the result reporting pipeline in various ways.
    tags = {
        'board': board,
        'build': build,
        # Required for proper rendering of MILO UI.
        'luci_project': 'chromeos',
        'skylab': 'run_suite',
        'skylab': 'staging',
        'suite': suite_name,
        'task_name': task_name,
    }
    osutils = autotest.chromite_load('osutils')
    with osutils.TempDir() as tempdir:
      summary_file = os.path.join(tempdir, 'summary.json')
      cmd = self._base_cmd('run') + [
          '--dimension', 'pool',  _SWARMING_POOL_SKYLAB_SUITE_BOTS,
          '--expiration', timeout_s,
          '--io-timeout', timeout_s,
          '--hard-timeout', timeout_s,
          '--print-status-update',
          '--priority', str(_TEST_PUSH_SUITE_PRIORITY),
          '--raw-cmd',
          '--task-name', task_name,
          '--task-summary-json', summary_file,
          '--timeout', timeout_s,
      ]
      for key, val in tags.iteritems():
        cmd += ['--tags', '%s:%s' % (key, val)]
      cmd += ['--'] + raw_cmd

      cros_build_lib = autotest.chromite_load('cros_build_lib')
      cros_build_lib.RunCommand(cmd, error_code_ok=True)
      return _extract_run_id(summary_file)


  def _base_cmd(self, subcommand):
    cmd = [
        self._cli_path, subcommand,
        '--swarming', self._host,
    ]
    if self._service_account_json is not None:
      cmd += ['--auth-service-account-json', self._service_account_json]
    return cmd


  def task_url(self, task_id):
    """Generate the task url based on task id."""
    return '%s/user/task/%s' % (self._host, task_id)


def _extract_run_id(path):
  if not os.path.isfile(path):
    raise errors.TestPushError('No task summary at %s' % path)
  with open(path) as f:
    summary = json.load(f)
  if not summary.get('shards') or len(summary['shards']) != 1:
    raise errors.TestPushError('Corrupted task summary at %s' % path)
  run_id = summary['shards'][0].get('run_id')
  if not run_id:
    raise errors.TestPushError('No run_id in task summary at %s' % path)
  return run_id


def _old_style_pool_label(label):
  _POOL_LABEL_PREFIX = 'dut_pool_'
  label = label.lower()
  if label.startswith(_POOL_LABEL_PREFIX):
    return label[len(_POOL_LABEL_PREFIX):]
  return label