# -*- coding: utf-8 -*-
# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Module to use goma from buildbot."""

from __future__ import print_function

import collections
import datetime
import getpass
import glob
import json
import os
import shlex
import sys
import tempfile

from chromite.lib import cros_build_lib
from chromite.lib import cros_logging as logging
from chromite.lib import gs
from chromite.lib import osutils
from chromite.lib import path_util


assert sys.version_info >= (3, 6), 'This module requires Python 3.6+'


_GOMA_COMPILER_PROXY_LOG_URL_TEMPLATE = (
    'https://chromium-build-stats.appspot.com/compiler_proxy_log/%s/%s')
_GOMA_NINJA_LOG_URL_TEMPLATE = (
    'https://chromium-build-stats.appspot.com/ninja_log/%s/%s')
GomaApproach = collections.namedtuple(
    'GomaApproach', ['rpc_extra_params',
                     'server_host',
                     'arbitrary_toolchain_support'])


# TODO(crbug.com/1035114) Refactor.
class Goma(object):
  """Interface to use goma on bots."""

  # Default environment variables to use goma.
  _DEFAULT_ENV_VARS = {
      # Set MAX_COMPILER_DISABLED_TASKS to let goma enter Burst mode, if
      # there are too many local fallback failures. In the Burst mode, goma
      # tries to use CPU cores as many as possible. Note that, by default,
      # goma runs only a few local fallback tasks in parallel at once.
      # The value is the threashold of the number of local fallback failures
      # to enter the mode.
      # Note that 30 is just heuristically chosen by discussion with goma team.
      #
      # Specifically, this is short-term workaround of the case that all
      # compile processes get local fallback. Otherwise, because goma uses only
      # several processes for local fallback by default, it causes significant
      # slow down of the build.
      # Practically, this happens when toolchain is updated in repository,
      # but prebuilt package is not yet ready. (cf. crbug.com/728971)
      'GOMA_MAX_COMPILER_DISABLED_TASKS': '30',

      # Disable goma soft stickiness.
      # Goma was historically using `soft stickiness cookie` so that uploaded
      # file cache is available as much as possible. However, such sticky
      # requests are cuasing unbalanced server load, and the disadvantage of the
      # unbalanceness cannot be negligible now. According to chrome's
      # experiment, the disadvantage of disabling soft stickiness is negligible,
      # and achieving balanced server load will have more advantage for entire
      # build. (cf. crbug.com/730962)
      # TODO(shinyak): This will be removed after crbug.com/730962 is resolved.
      'GOMA_BACKEND_SOFT_STICKINESS': 'false',

      # Enable DepsCache. DepsCache is a cache that holds a file list that
      # compiler_proxy sends to goma server for each compile. This can
      # reduces a lot of I/O and calculation.
      # This is the base file name under GOMA_CACHE_DIR.
      'GOMA_DEPS_CACHE_FILE': 'goma.deps',

      # Only run one command in parallel per core.
      #
      # TODO(crbug.com/998076): Increase if Goma fork issue is fixed.
      'NINJA_CORE_MULTIPLIER': '1',
  }

  def __init__(self, goma_dir, goma_client_json, goma_tmp_dir=None,
               stage_name=None, chromeos_goma_dir=None, chroot_dir=None,
               goma_approach=None, log_dir=None, stats_filename=None,
               counterz_filename=None):
    """Initializes Goma instance.

    This ensures that |self.goma_log_dir| directory exists (if missing,
    creates it).

    Args:
      goma_dir: Path to the Goma client used for simplechrome
                (outside of chroot).
      goma_client_json: Path to the service account json file to use goma.
        On bots, this must be specified, otherwise raise a ValueError.
        On local, this is optional, and can be set to None.
      goma_tmp_dir: Path to the GOMA_TMP_DIR to be passed to goma programs.
        If given, it is used. If not given, creates a directory under
        /tmp in the chroot, expecting that the directory is removed in the
        next run's clean up phase on bots.
      stage_name: optional name of the currently running stage. E.g.
        "build_packages" or "test_simple_chrome_workflow". If this is set
        deps cache is enabled.
      chromeos_goma_dir: Path to the Goma client used for build package.
                         path should be represented as outside of chroot.
                         If None, goma_dir will be used instead.
      chroot_dir: The base chroot path to use when the chroot path is not at
        the default location.
      goma_approach: Indicates some extra environment variables to set when
        testing alternative goma approaches.
      log_dir: Allows explicitly setting the log directory. Used for the
        Build API for extracting the logs afterwords. Should be the log
        directory inside the chroot, based on the chroot path when
        outside the chroot.
      stats_filename: The name of the file to use for the GOMA_DUMP_STATS_FILE
        setting. The file will be created in the log directory.
      counterz_filename: The name of the file to use for the
        GOMA_DUMP_COUNTERZ_FILE setting. The file will be created in the log
        directory.

    Raises:
      ValueError if 1) |goma_dir| does not point to a directory, 2)
      on bots, but |goma_client_json| is not given, 3) |goma_client_json|
      is given, but it does not point to a file, or 4) if |goma_tmp_dir| is
      given but it does not point to a directory.
    """
    # Sanity checks of given paths.
    if not os.path.isdir(goma_dir):
      raise ValueError('goma_dir does not point a directory: %s' % (goma_dir,))

    # If this script runs on bot, service account json file needs to be
    # provided, otherwise it cannot access to goma service.
    if cros_build_lib.HostIsCIBuilder() and goma_client_json is None:
      raise ValueError(
          'goma is enabled on bot, but goma_client_json is not provided')

    # If goma_client_json file is provided, it must be an existing file.
    if goma_client_json and not os.path.isfile(goma_client_json):
      raise ValueError(
          'Goma client json file is missing: %s' % (goma_client_json,))

    # If goma_tmp_dir is provided, it must be an existing directory.
    if goma_tmp_dir and not os.path.isdir(goma_tmp_dir):
      raise ValueError(
          'GOMA_TMP_DIR does not point a directory: %s' % (goma_tmp_dir,))

    self.linux_goma_dir = goma_dir
    self.chromeos_goma_dir = chromeos_goma_dir
    self.goma_approach = goma_approach
    # If Goma dir for ChromeOS SDK does not set, fallback to use goma_dir.
    if self.chromeos_goma_dir is None:
      self.chromeos_goma_dir = goma_dir
    # Sanity checks of given paths.
    if not os.path.isdir(self.chromeos_goma_dir):
      raise ValueError('chromeos_goma_dir does not point a directory: %s' % (
          self.chromeos_goma_dir,))

    self.goma_client_json = goma_client_json
    if stage_name:
      self.goma_cache = os.path.join(goma_dir, 'goma_cache', stage_name)
      osutils.SafeMakedirs(self.goma_cache)
    else:
      self.goma_cache = None

    if goma_tmp_dir is None:
      # path_util depends on the chroot directory existing at
      # SOURCE_ROOT/chroot. This assumption is not valid for Luci builders,
      # but generally shouldn't be an assumption anyway since we allow setting
      # the chroot location. This block, and the two a few lines down, bypass
      # path_util to compensate for those assumptions.
      # TODO(crbug.com/1014138) Cleanup when path_util can handle custom chroot.
      if chroot_dir:
        chroot_tmp = os.path.join(chroot_dir, 'tmp')
      else:
        chroot_tmp = path_util.FromChrootPath('/tmp')

      # If |goma_tmp_dir| is not given, create GOMA_TMP_DIR (goma
      # compiler_proxy's working directory), and its log directory.
      # Create unique directory by mkdtemp under chroot's /tmp.
      # Expect this directory is removed in next run's clean up phase.
      goma_tmp_dir = tempfile.mkdtemp(
          prefix='goma_tmp_dir.',
          dir=chroot_tmp)
    self.goma_tmp_dir = goma_tmp_dir

    if chroot_dir:
      self.chroot_goma_tmp_dir = os.path.join(
          '/', os.path.relpath(self.goma_tmp_dir, chroot_dir))
    else:
      self.chroot_goma_tmp_dir = path_util.ToChrootPath(self.goma_tmp_dir)

    self._log_dir = log_dir
    # Create log directory if not exist.
    if not os.path.isdir(self.goma_log_dir):
      os.mkdir(self.goma_log_dir)

    self._stats_file = None
    self._chroot_stats_file = None
    self._counterz_file = None
    self._chroot_counterz_file = None
    if stats_filename:
      self._stats_file = os.path.join(self.goma_log_dir, stats_filename)
    if counterz_filename:
      self._counterz_file = os.path.join(self.goma_log_dir, counterz_filename)

    if chroot_dir:
      self.chroot_goma_log_dir = os.path.join(
          '/', os.path.relpath(self.goma_log_dir, chroot_dir))
      if self._stats_file:
        self._chroot_stats_file = os.path.join(
            '/', os.path.relpath(self._stats_file, chroot_dir))
      if self._counterz_file:
        self._chroot_counterz_file = os.path.join(
            '/', os.path.relpath(self._counterz_file, chroot_dir))
    else:
      self.chroot_goma_log_dir = path_util.ToChrootPath(self.goma_log_dir)
      if self._stats_file:
        self._chroot_stats_file = path_util.ToChrootPath(self._stats_file)
      if self._counterz_file:
        self._chroot_counterz_file = path_util.ToChrootPath(self._counterz_file)

  @property
  def goma_log_dir(self):
    """Path to goma's log directory."""
    return self._log_dir or os.path.join(self.goma_tmp_dir, 'log_dir')

  def GetExtraEnv(self):
    """Extra env vars set to use goma."""
    result = dict(
        Goma._DEFAULT_ENV_VARS,
        GOMA_DIR=self.linux_goma_dir,
        GOMA_TMP_DIR=self.goma_tmp_dir,
        GLOG_log_dir=self.goma_log_dir)

    self._AddCommonExtraEnv(result)

    if self.goma_client_json:
      result['GOMA_SERVICE_ACCOUNT_JSON_FILE'] = self.goma_client_json

    if self.goma_cache:
      result['GOMA_CACHE_DIR'] = self.goma_cache

    if self._stats_file:
      result['GOMA_DUMP_STATS_FILE'] = self._stats_file
    if self._counterz_file:
      result['GOMA_DUMP_COUNTERZ_FILE'] = self._counterz_file
      result['GOMA_ENABLE_COUNTERZ'] = 'true'

    return result

  def GetChrootExtraEnv(self):
    """Extra env vars set to use goma inside chroot."""
    # Note: GOMA_DIR and GOMA_SERVICE_ACCOUNT_JSON_FILE in chroot is hardcoded.
    # Please see also enter_chroot.sh.
    goma_dir = os.path.join('/home', os.environ.get('USER'), 'goma')
    result = dict(
        Goma._DEFAULT_ENV_VARS,
        GOMA_DIR=goma_dir,
        GOMA_TMP_DIR=self.chroot_goma_tmp_dir,
        GLOG_log_dir=self.chroot_goma_log_dir)

    self._AddCommonExtraEnv(result)

    if self.goma_client_json:
      result['GOMA_SERVICE_ACCOUNT_JSON_FILE'] = (
          '/creds/service_accounts/service-account-goma-client.json')

    if self.goma_cache:
      result['GOMA_CACHE_DIR'] = os.path.join(
          goma_dir, os.path.relpath(self.goma_cache, self.chromeos_goma_dir))

    if self._chroot_stats_file:
      result['GOMA_DUMP_STATS_FILE'] = self._chroot_stats_file
    if self._chroot_counterz_file:
      result['GOMA_DUMP_COUNTERZ_FILE'] = self._chroot_counterz_file
      result['GOMA_ENABLE_COUNTERZ'] = 'true'

    return result

  def _AddCommonExtraEnv(self, result):
    """Sets extra env vars to use goma common to in / out chroot."""
    if self.goma_approach:
      result['GOMA_RPC_EXTRA_PARAMS'] = self.goma_approach.rpc_extra_params
      result['GOMA_SERVER_HOST'] = self.goma_approach.server_host
      result['GOMA_ARBITRARY_TOOLCHAIN_SUPPORT'] = (
          'true' if self.goma_approach.arbitrary_toolchain_support else 'false')

  def _RunGomaCtl(self, command):
    goma_ctl = os.path.join(self.linux_goma_dir, 'goma_ctl.py')
    # TODO(crbug.com/1007384): Stop forcing Python 2.
    cros_build_lib.run(
        ['python2', goma_ctl, command], extra_env=self.GetExtraEnv())

  def Start(self):
    """Starts goma compiler proxy."""
    self._RunGomaCtl('start')

  def Stop(self):
    """Stops goma compiler proxy."""
    self._RunGomaCtl('stop')

  def UploadLogs(self, cbb_config_name):
    """Uploads INFO files related to goma.

    Args:
      cbb_config_name: name of cbb_config.

    Returns:
      URL to the compiler_proxy log visualizer. None if unavailable.
    """
    uploader = GomaLogUploader(self.goma_log_dir,
                               cbb_config_name=cbb_config_name)
    return uploader.Upload()


# Note: Public for testing purpose. In real use, please think about using
# Goma.UploadLogs() instead.
class GomaLogUploader(object):
  """Manages to upload goma log files."""

  # The Google Cloud Storage bucket to store logs related to goma.
  _BUCKET = 'chrome-goma-log'

  def __init__(self, goma_log_dir, today=None, dry_run=False,
               cbb_config_name=''):
    """Initializes the uploader.

    Args:
      goma_log_dir: path to the directory containing goma's INFO log files.
      today: datetime.date instance representing today. This is for testing
        purpose, because datetime.date is unpatchable. In real use case,
        this must be None.
      dry_run: If True, no actual upload. This is for testing purpose.
      cbb_config_name: Name of cbb_config.
    """
    self._goma_log_dir = goma_log_dir
    logging.info('Goma log directory is: %s', self._goma_log_dir)

    # Set log upload destination.
    if today is None:
      today = datetime.date.today()
    self.dest_path = '%s/%s' % (
        today.strftime('%Y/%m/%d'), cros_build_lib.GetHostName())
    self._remote_dir = 'gs://%s/%s' % (GomaLogUploader._BUCKET, self.dest_path)
    logging.info('Goma log upload destination: %s', self._remote_dir)

    # HACK(yyanagisawa): I suppose LUCI do not set BUILDBOT_BUILDERNAME.
    is_luci = not bool(os.environ.get('BUILDBOT_BUILDERNAME'))
    # Build metadata to be annotated to log files.
    # Use OrderedDict for json output stabilization.
    builder_info = collections.OrderedDict([
        ('builder', os.environ.get('BUILDBOT_BUILDERNAME', '')),
        ('master', os.environ.get('BUILDBOT_MASTERNAME', '')),
        ('slave', os.environ.get('BUILDBOT_SLAVENAME', '')),
        ('clobber', bool(os.environ.get('BUILDBOT_CLOBBER'))),
        ('os', 'chromeos'),
        ('is_luci', is_luci),
        ('cbb_config_name', cbb_config_name),
    ])
    if is_luci:
      # TODO(yyanagisawa): will adjust to valid value if needed.
      builder_info['builder_id'] = collections.OrderedDict([
          ('project', 'chromeos'),
          ('builder', 'Prod'),
          ('bucket', 'general'),
      ])
    builder_info_json = json.dumps(builder_info)
    logging.info('BuilderInfo: %s', builder_info_json)
    self._headers = ['x-goog-meta-builderinfo:' + builder_info_json]

    self._gs_context = gs.GSContext(dry_run=dry_run)

  def Upload(self):
    """Uploads all necessary log files to Google Storage.

    Returns:
      A list of pairs of label and URL of goma log visualizers to be linked
      from the build status page.
    """
    compiler_proxy_subproc_paths = self._UploadInfoFiles(
        'compiler_proxy-subproc')
    # compiler_proxy-subproc.INFO file should be exact one.
    if len(compiler_proxy_subproc_paths) != 1:
      logging.warning('Unexpected compiler_proxy-subproc INFO files: %r',
                      compiler_proxy_subproc_paths)

    compiler_proxy_paths = self._UploadInfoFiles('compiler_proxy')
    # compiler_proxy.INFO file should be exact one.
    if len(compiler_proxy_paths) != 1:
      logging.warning('Unexpected compiler_proxy INFO files: %r',
                      compiler_proxy_paths)
    compiler_proxy_path, uploaded_compiler_proxy_filename = (
        compiler_proxy_paths[0] if compiler_proxy_paths else (None, None))

    self._UploadGomaccInfoFiles()

    uploaded_ninja_log_filename = self._UploadNinjaLog(compiler_proxy_path)

    # Build URL to be linked.
    result = []
    if uploaded_compiler_proxy_filename:
      result.append((
          'Goma compiler_proxy log',
          _GOMA_COMPILER_PROXY_LOG_URL_TEMPLATE % (
              self.dest_path, uploaded_compiler_proxy_filename)))
    if uploaded_ninja_log_filename:
      result.append((
          'Goma ninja_log',
          _GOMA_NINJA_LOG_URL_TEMPLATE % (
              self.dest_path, uploaded_ninja_log_filename)))
    return result

  def _UploadInfoFiles(self, pattern):
    """Uploads INFO files matched with pattern, with gzip'ing.

    Args:
      pattern: matching path pattern.

    Returns:
      A list of uploaded file paths.
    """
    # Find files matched with the pattern in |goma_log_dir|. Sort for
    # stabilization.
    paths = sorted(glob.glob(
        os.path.join(self._goma_log_dir, '%s.*.INFO.*' % pattern)))
    if not paths:
      logging.warning('No glog files matched with: %s', pattern)

    result = []
    for path in paths:
      logging.info('Uploading %s', path)
      uploaded_filename = os.path.basename(path) + '.gz'
      self._gs_context.CopyInto(
          path, self._remote_dir, filename=uploaded_filename,
          auto_compress=True, headers=self._headers)
      result.append((path, uploaded_filename))
    return result

  def _UploadGomaccInfoFiles(self):
    """Uploads gomacc INFO files, with gzip'ing.

    Returns:
      Uploaded file path. If failed, None.
    """

    # Since the number of gomacc logs can be large, we'd like to compress them.
    # Otherwise, upload will take long (> 10 mins).
    # Each gomacc logs file size must be small (around 4KB).

    # Find files matched with the pattern in |goma_log_dir|.
    # The paths were themselves used as the inputs for the create
    # tarball, but there can be too many of them. As long as we have
    # files we'll just tar up the entire directory.
    gomacc_paths = glob.glob(os.path.join(self._goma_log_dir,
                                          'gomacc.*.INFO.*'))
    if not gomacc_paths:
      # gomacc logs won't be made every time.
      # Only when goma compiler_proxy has
      # crashed. So it's usual gomacc logs are not found.
      logging.info('No gomacc logs found')
      return None

    # Taking the alphabetically first name as uploaded_filename.
    tgz_name = os.path.basename(min(gomacc_paths)) + '.tar.gz'
    # When using the pigz compressor (what we use for gzip) to create an
    # archive in a folder that is also a source for contents, there is a race
    # condition involving the created archive itself that can cause it to fail
    # creating the archive. To avoid this, make the archive in a tempdir.
    with osutils.TempDir() as tempdir:
      tgz_path = os.path.join(tempdir, tgz_name)
      cros_build_lib.CreateTarball(target=tgz_path,
                                   cwd=self._goma_log_dir,
                                   compression=cros_build_lib.COMP_GZIP)
      self._gs_context.CopyInto(tgz_path, self._remote_dir,
                                filename=tgz_name,
                                headers=self._headers)
    return tgz_name

  def _UploadNinjaLog(self, compiler_proxy_path):
    """Uploads .ninja_log file and its related metadata.

    This uploads the .ninja_log file generated by ninja to build Chrome.
    Also, it appends some related metadata at the end of the file following
    '# end of ninja log' marker.

    Args:
      compiler_proxy_path: Path to the compiler proxy, which will be contained
        in the metadata.

    Returns:
      The name of the uploaded file.
    """
    ninja_log_path = os.path.join(self._goma_log_dir, 'ninja_log')
    if not os.path.exists(ninja_log_path):
      logging.warning('ninja_log is not found: %s', ninja_log_path)
      return None
    ninja_log_content = osutils.ReadFile(ninja_log_path)

    try:
      st = os.stat(ninja_log_path)
      ninja_log_mtime = datetime.datetime.fromtimestamp(st.st_mtime)
    except OSError:
      logging.exception('Failed to get timestamp: %s', ninja_log_path)
      return None

    ninja_log_info = self._BuildNinjaInfo(compiler_proxy_path)

    # Append metadata at the end of the log content.
    ninja_log_content += '# end of ninja log\n' + json.dumps(ninja_log_info)

    # Aligned with goma_utils in chromium bot.
    pid = os.getpid()

    upload_ninja_log_path = os.path.join(
        self._goma_log_dir,
        'ninja_log.%s.%s.%s.%d' % (
            getpass.getuser(), cros_build_lib.GetHostName(),
            ninja_log_mtime.strftime('%Y%m%d-%H%M%S'), pid))
    osutils.WriteFile(upload_ninja_log_path, ninja_log_content)
    uploaded_filename = os.path.basename(upload_ninja_log_path) + '.gz'
    self._gs_context.CopyInto(
        upload_ninja_log_path, self._remote_dir, filename=uploaded_filename,
        auto_compress=True, headers=self._headers)
    return uploaded_filename

  def _BuildNinjaInfo(self, compiler_proxy_path):
    """Reads metadata for the ninja run.

    Each metadata should be written into a dedicated file in the log directory.
    Read the info, and build the dict containing metadata.

    Args:
      compiler_proxy_path: Path to the compiler_proxy log file.

    Returns:
      A dict of the metadata.
    """

    info = {'platform': 'chromeos'}

    command_path = os.path.join(self._goma_log_dir, 'ninja_command')
    if os.path.exists(command_path):
      info['cmdline'] = shlex.split(
          osutils.ReadFile(command_path).strip())

    cwd_path = os.path.join(self._goma_log_dir, 'ninja_cwd')
    if os.path.exists(cwd_path):
      info['cwd'] = osutils.ReadFile(cwd_path).strip()

    exit_path = os.path.join(self._goma_log_dir, 'ninja_exit')
    if os.path.exists(exit_path):
      info['exit'] = int(osutils.ReadFile(exit_path).strip())

    env_path = os.path.join(self._goma_log_dir, 'ninja_env')
    if os.path.exists(env_path):
      # env is null-byte separated, and has a trailing null byte.
      content = osutils.ReadFile(env_path).rstrip('\0')
      info['env'] = dict(line.split('=', 1) for line in content.split('\0'))

    if compiler_proxy_path:
      info['compiler_proxy_info'] = os.path.basename(compiler_proxy_path)

    return info