# 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 urllib

from lucifer import autotest

# This is hard-coded everywhere -- on builders requesting skylab suite as well
# as in commands to be generated by users requesting one off suites.
_SWARMING_POOL_SKYLAB_BOTS = 'ChromeOSSkylab'


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 _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