"""Build step class for building OTA's."""

from __future__ import absolute_import
import logging
import os
import re
import sys

from helpers import make_utils
from helpers import ninja_utils
from helpers import py_env_utils
from slave import base_step

# Taken from recommended value in autoninja
BUILD_ACCELERATION_JOB_MULTIPLIER = 40

# Allow sufficient time for builds without build acceleration to finish.
_OTA_STEP_MAX_TIME_SECONDS = 6 * 60 * 60

_MESH_TAG = '-mesh-'
_YETI_TAG = '-yeti-'
_SDM_TAG = '-sdm-'
_GALLIUM_TAG = '-gallium-'
_OPAL_TAG = '-opal-'
_CONCIERGE_TAG = '-concierge-'
_HEALTH_TAG = '-health-'
_FUCHSIA_TAG = '-fuchsia-'
_FUCHSIA_REV_TAG = '-fuchsia_rev-'
_ML_FRAMEWORK_TAG = '-ml_framework-'
_DEMO_TAG = '-demo-'
_OTA_FILES_RE = r'.*-ota-.*\..*'
_PARTNER_TAG = '-partner-'
_LEGACY_TAG = '_legacy'
_1LED_TAG = '-1led-'
_MODULAR_COMMS_TAG = '-modular_comms-'
_IOT_CAST_TAG = '-iot_cast-'

GANETI_TAG = '-ganeti-'

_OTA_ZIP_TEMPLATE = '{}-ota-{}.zip'


class OtaStep(base_step.BaseStep):
  """Build step class for building OTA's."""

  def __init__(self, build_name, name='build ota', board_name=None,
               extra_params=None, max_ota_size=None, target_name=None,
               timeout_secs=_OTA_STEP_MAX_TIME_SECONDS,
               **kwargs):
    """Creates a OtaStep instance.

    Args:
      build_name: Name of build (e.g. anchovy-eng, tvdefault-eng)
      name: user-visible name of this step.
      board_name: optional BOARD_NAME param for make if any.
      extra_params: optional list of params to add
      max_ota_size: Max size of OTA to enfoce if any enforcement required.
      timeout_secs: Max execution time for OTA steps.
      **kwargs: Any additional args to pass to BaseStep.
    """
    base_step.BaseStep.__init__(self, name=name, timeout_secs=timeout_secs,
                                **kwargs)
    self._build_name = build_name
    self._board_name = board_name
    self._extra_params = extra_params
    self._max_ota_size = max_ota_size
    self._target_name = target_name

    self._enable_verbose_logging = self._properties.get(
      'enable_verbose_logging', False)

    self._is_mesh_build = _MESH_TAG in self._build_name
    self._build_name = self._build_name.replace(_MESH_TAG, '-')

    self._is_yeti_build = _YETI_TAG in self._build_name
    self._build_name = self._build_name.replace(_YETI_TAG, '-')

    self._is_sdm_build = _SDM_TAG in self._build_name
    self._build_name = self._build_name.replace(_SDM_TAG, '-')

    self._is_gallium_build = _GALLIUM_TAG in self._build_name
    self._build_name = self._build_name.replace(_GALLIUM_TAG, '-')

    self._is_opal_build = _OPAL_TAG in self._build_name
    self._build_name = self._build_name.replace(_OPAL_TAG, '-')

    self._is_concierge_build = _CONCIERGE_TAG in self._build_name
    self._build_name = self._build_name.replace(_CONCIERGE_TAG, '-')

    self._is_health_build = _HEALTH_TAG in self._build_name
    self._build_name = self._build_name.replace(_HEALTH_TAG, '-')

    self._is_fuchsia_build = _FUCHSIA_TAG in self._build_name
    self._build_name = self._build_name.replace(_FUCHSIA_TAG, '-')

    self._is_fuchsia_rev_build = _FUCHSIA_REV_TAG in self._build_name
    self._build_name = self._build_name.replace(_FUCHSIA_REV_TAG, '-')

    self._is_ml_framework_build = _ML_FRAMEWORK_TAG in self._build_name
    self._build_name = self._build_name.replace(_ML_FRAMEWORK_TAG, '-')

    self._is_demo_build = _DEMO_TAG in self._build_name
    self._build_name = self._build_name.replace(_DEMO_TAG, '-')

    self._is_ganeti_build = GANETI_TAG in self._build_name
    self._build_name = self._build_name.replace(GANETI_TAG, '-')

    self._is_1led_build = _1LED_TAG in self._build_name
    self._build_name = self._build_name.replace(_1LED_TAG, '-')

    self._is_modular_comms_build = _MODULAR_COMMS_TAG in self._build_name
    self._build_name = self._build_name.replace(_MODULAR_COMMS_TAG, '-')

    self._is_iot_cast_build = _IOT_CAST_TAG in self._build_name
    self._build_name = self._build_name.replace(_IOT_CAST_TAG, '-')

    # Partner repo is specified in properties, but the tag needs to be
    # removed from the build name.
    self._build_name = self._build_name.replace(_PARTNER_TAG, '-')

    self._build_name = self._build_name.replace(_LEGACY_TAG, '')

  def _dist_dir(self):
    dist_dir = os.path.join('out/dist', self.build_number)
    if self.patchset != '0':
      dist_dir = os.path.join(dist_dir, str(self.patchset))
    return dist_dir

  def make_target_param(self, target_params=False):
    """Creates PRODUCT or TARGET params for make.

    build/build test steps are using PRODUCT-<build name> notation,
    TARGET_PRODUCT and TARGET_BUILD_VARIANT params required for OSS and FCT.

    Args:
      target_params: use TARGET style params regardless current build_system.
    Returns:
      array of params to be appended to make command.
    """
    target_name = self._target_name
    if not target_name:
      target_name = self._build_name
    if target_params:
      product, variant = target_name.split('-')
      return ['TARGET_PRODUCT={}'.format(product),
              'TARGET_BUILD_VARIANT={}'.format(variant)]
    return ['PRODUCT-' + target_name]

  def make_command_env(self):
    """Generates env dictionary for make command."""
    return self.build_accelerator.environment_variables

  def _make_extra_flags(self):
    """Generates make extra params based on build type."""

    flags = []
    variant = self._build_name.split('-')[-1]
    if variant != 'eng':
      flags += ['OFFICIAL_BUILD=1']
    if self._enable_verbose_logging:
      flags.append('showcommands')
    if self._is_mesh_build:
      flags += ['USE_GWIFI=1']
    if self._is_yeti_build:
      flags += ['ENABLE_YETI=1', 'ENABLE_CLASSIC_YETI=0']
    if self._is_sdm_build:
      flags += ['ENABLE_HOSPITALITY=1', 'ENABLE_MANAGED_MODE=1']
    if self._is_gallium_build:
      flags += ['ENABLE_GALLIUM=1']
    if self._is_opal_build:
      flags += ['ENABLE_OPAL=true']
    if self._is_concierge_build:
      flags += ['ENABLE_SOUND_PRESENCE=1']
    if self._is_health_build:
      flags += ['ENABLE_HEALTH_SOUND=1', 'ENABLE_MANAGED_MODE=1']
    if self._is_fuchsia_build:
      flags += ['FUCHSIA_OTA=true']
    if self._is_fuchsia_rev_build:
      flags += ['FUCHSIA_REVERSE=true']
    if self._is_ml_framework_build:
      flags += ['ENABLE_ML_FRAMEWORK=true']
    if self._is_demo_build:
      flags += ['DEMO_OTA=true']
    if self._board_name:
      flags += ['BOARD_NAME={}'.format(self._board_name)]
    if self._is_1led_build:
      flags += ['NUM_LEDS=1']
    if self._is_modular_comms_build:
      flags += ['ENABLE_MODULAR_COMMS=true']
    if self._is_iot_cast_build:
      flags += ['ENABLE_IOT_CAST=1']
    if self._extra_params:
      flags += self._extra_params
    return flags

  def _build_ota(self, custom_shell):
    """Builds |build_name| OTA package for the |issue| and |patchset|.

    Args:
      custom_shell: List containing custom shell, to be used with make.
    Returns:
      True if the OTA package was built successfully, False otherwise.
    """
    # Clean the systemimage first to avoid bleedover on incremental builds
    # from previous builds of the same product.
    command = ['make',
               'BUILD_NUMBER={}'.format(self.build_number),
               'USER=auto',
               'clean-systemimage',
               'DIST_DIR={}'.format(self._dist_dir())]
    command += self.build_accelerator.make_flags

    quiet = False
    if custom_shell:
      quiet = True

    # For this step, don't use PRODUCT-* because it adds a "dist" goal. Just
    # perform the clean-systemimage.
    command += self.make_target_param(target_params=True)
    command += self._make_extra_flags()
    returncode, stdout, stderr = self.exec_subprocess(
        command, env=self.make_command_env(), quiet=quiet)

    # We permit this step to fail because the clean-systemimage step
    # is custom to Cast OS and does not exist on older release branches.
    # Since CQ build logic always runs on master, regardless of the age
    # of the code being built, we must maintain backwards compatibility.
    if returncode != 0:
      logging.warning(
        'Failed to run clean-systemimage. This is normal on release branches '
        'from 1.50 and earlier. If this happens on a newer branch it is an '
        'error and will cause non-deterministic ota package sizes.\n'
        'Return code %d, stderr:\n%s',
        returncode, stderr)


    command = ['make',
               '--output-sync=target',
               '-j{}'.format(self.get_num_jobs(
                   multiplier=BUILD_ACCELERATION_JOB_MULTIPLIER)),
               'BUILD_NUMBER={}'.format(self.build_number),
               'USER=auto',
               'dist',
               'DIST_DIR={}'.format(self._dist_dir()),
               'DISABLE_AUTO_INSTALLCLEAN=true',]
    command += self.build_accelerator.make_flags
    command += self.make_target_param()
    command += self._make_extra_flags()

    if custom_shell:
      command += custom_shell

    returncode, stdout, stderr = self.exec_subprocess(
        command, env=self.make_command_env(), quiet=quiet)

    if custom_shell:
      self._merge_build_logs()

    if returncode != 0:
      if custom_shell:
        self._dump_error_logs(stdout, stderr)
      self._error_processor(stdout, stderr)
    else:
      if custom_shell:
        self._print_build_success()

    if custom_shell:
      sys.stdout.write('Complete logs can be found in build_ota.log '
      'within invocation artifacts, sorted by module. Logs that are '
      'generated outside of a module are logged under the "general" module.\n')
      sys.stdout.flush()

    return returncode == 0

  def _error_processor(self, stdout, stderr):
    messages = ninja_utils.ninja_failure_to_comments(stdout, max_errors=3)
    for message in messages:
      self.add_review({'message': message})

    comments = make_utils.make_error_to_comments(stderr, self.directory)
    if comments:
      self.add_review({'comments': comments})

  def _print_build_success(self):
    sys.stdout.write('Build successful.\n')
    sys.stdout.flush()

  def _clear_dist(self, quiet):
    """Clears everything in the dist directory leftover from previous builds.
    Args:
      quiet: Suppress subprocess output.
    """
    dist_dirs = self.glob_files('out/dist/*')
    if dist_dirs:
      self.exec_subprocess(['rm', '-rf'] + dist_dirs, quiet=quiet)

  def _clear_ota_artifacts(self, quiet):
    """Removes all OTA artifacts leftover by previous builds.
    Args:
      quiet: Suppress subprocess output.
    """
    archive_pattern = re.compile(_OTA_FILES_RE)

    dist_files = self.glob_files('{}/*'.format(self._dist_dir()))
    ota_files = [f for f in dist_files if archive_pattern.match(f)]
    if ota_files:
      self.exec_subprocess(['rm', '-f'] + ota_files, quiet=quiet)

  def _clear_symbols_artifacts(self, quiet):
    """Removes all symbols artifacts leftover by previous builds.
    Args:
      quiet: Suppress subprocess output.
    """
    symbols_files = self.glob_files('out/target/product/*/*-symbols-*.zip')
    if symbols_files:
      self.exec_subprocess(['rm', '-f'] + symbols_files, quiet=quiet)

  def _ota_archive_output_dir(self):
    """Return system specific archive dir."""
    output_dir = self.get_gcs_dir()
    if self.build_system != 'catabuilder':
      output_dir = os.path.join(output_dir, 'artifacts')
      if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    return output_dir

  def _archive_artifacts(self, quiet):
    """Copies the artifacts for archival.
    Args:
      quiet: Suppress subprocess output.
    """
    output_dir = self._ota_archive_output_dir()
    if self.build_system == 'catabuilder':
      # Archive everything in the continuous build system.
      archive_pattern = re.compile(r'.*')
    else:
      archive_pattern = re.compile(_OTA_FILES_RE)

    dist_files = self.glob_files('{}/*'.format(self._dist_dir()))
    archive_files = [f for f in dist_files if archive_pattern.match(f)]
    for f in archive_files:
      self.exec_subprocess(['cp', '-r', f, output_dir], quiet=quiet)

    # Save the OTA file list for later use.
    ota_files = [f for f in archive_files if re.match(_OTA_FILES_RE, f)]
    if ota_files:
      filenames = [os.path.basename(x) for x in ota_files]
      self.set_build_property('ota_files', filenames)

  def _dump_error_logs(self, stdout, stderr):
    """Dumps error logs to stderr."""
    product = self._build_name.split('-')[0]
    error_log_files = sorted(self.glob_files(
      'out/target/product/' + product + '/build_logs/*.err.txt'),
      key=os.path.basename)

    if error_log_files:
      sys.stdout.write('The following modules failed to build:\n')
      sys.stdout.flush()
      for error_file in error_log_files:
        with open(error_file, 'r') as err_file:
          contents = err_file.read()
          sys.stdout.write(os.path.basename(error_file).split('.')[0] + ':\n')
          sys.stdout.flush()
          sys.stderr.write(contents)
          sys.stderr.flush()
    else:
      # No errors found, dump error logs to test.log.
      sys.stderr.write(stderr)
      sys.stderr.flush()

  def _merge_build_logs(self):
    """Copies build logs for archival."""
    product = self._build_name.split('-')[0]
    build_log_files = sorted(self.glob_files(
      'out/target/product/' + product + '/build_logs/*.out.txt'),
      key=os.path.basename)
    output_dir = self.get_gcs_dir()
    if not os.path.exists(output_dir):
      os.makedirs(output_dir)
    build_logs_file = os.path.join(output_dir, 'build_ota.log')
    with open(build_logs_file, 'w') as out_file:
      for log_file in build_log_files:
        with open(log_file, 'r') as in_file:
          contents = in_file.read()
          if contents:
            out_file.write('**********************************'
              '************************************\n');
            out_file.write(os.path.basename(log_file).split('.')[0] + ':\n')
            out_file.write('**********************************'
              '************************************\n');
            out_file.write(contents + '\n')
      out_file.flush()

  def _check_ota_size(self):
    """Checks the size of the OTA.

    Returns:
      Boolean if the OTA is under the max size.
    """
    # TODO(gfhuang): A/B ota size check need rework.
    # This as-is checks both .zip and .bin.
    archive_pattern = re.compile(_OTA_FILES_RE)

    dist_files = self.glob_files('{}/*'.format(self._dist_dir()))
    ota_files = [f for f in dist_files if archive_pattern.match(f)]
    self._log_ota_size(ota_files)
    if not self._max_ota_size:
      return True

    for ota_file in ota_files:
      ota_size = os.path.getsize(ota_file)
      if ota_size > self._max_ota_size:
        message = 'OTA size of {} exceed max of {} for file {}'.format(
            ota_size, self._max_ota_size, ota_file)
        sys.stdout.write(message + '\n')
        sys.stdout.flush()
        self.add_review({'message': message})
        return False
    return True

  def _log_ota_size(self, ota_files):
    target_ota_filename = _OTA_ZIP_TEMPLATE.format(
        self.get_property('buildername', default='').split('-')[0],
        self.get_property('buildset'))
    ota_file = next((f for f in ota_files
                     if os.path.basename(f) == target_ota_filename), None)

    if not ota_file:
      return

    self.set_build_property('ota_size', os.path.getsize(ota_file))
    if self._max_ota_size:
      self.set_build_property('max_ota_size', self._max_ota_size)

  def _make_custom_shell(self):
    """Returns a list containing a path to the custom shell"""
    shell_path = os.path.abspath(os.path.join(
                    os.getcwd(),
                    'build/android-mk-module-split-shell'))
    if os.path.isfile(shell_path):
      return ['ANDROID_BUILD_SHELL={}'.format(shell_path)]
    return []

  def run(self):
    """Builds ota package.

    Returns:
      True iff there were no errors.
    """
    custom_shell = self._make_custom_shell()
    self._clear_dist(bool(custom_shell))
    self._clear_ota_artifacts(bool(custom_shell))
    self._clear_symbols_artifacts(bool(custom_shell))
    if not self._build_ota(custom_shell):
      return False
    self._archive_artifacts(bool(custom_shell))
    if not self._check_ota_size():
      return False
    return True


class OtaPropStep(base_step.BaseStep):
  """Ota Property file step class for generating repo.prop."""

  def __init__(self, **kwargs):
    """Creates a OtaPropStep instance.

    Args:
      **kwargs: Any additional args to pass to BaseStep.
    """
    base_step.BaseStep.__init__(self, name='save repo.prop', **kwargs)

  def run(self):
    prop_cmd = [
        'repo', 'forall', '-c',
        'git show -s --format="$REPO_PROJECT %H"'
    ]
    proj_sha_pair_re = re.compile(r'^.+? [0-9a-f]{40}$', flags=re.IGNORECASE)

    repo_prop = self.exec_subprocess(prop_cmd, check_output=True, quiet=True)
    valid_lines = [l for l in repo_prop.split('\n')
                   if proj_sha_pair_re.match(l.strip())]
    prop_file = os.path.join(self.get_gcs_dir(), 'repo.prop')
    with open(prop_file, 'w') as f:
      f.write('\n'.join(valid_lines))
    return True

class OtaCameraTargetsStep(base_step.BaseStep):
  """Ota Property file step class for generating repo.prop."""

  def __init__(self, product_name, is_eng_build, name='camera targets check',
               **kwargs):
    """Creates a OtaCameraTargets instance.

    Args:
      product_name: product name without variant, e.g. spencer, venus
      is_eng_build: if this build is eng build
      name: user-visible name of this step.
      **kwargs: Any additional args to pass to BaseStep.
    """
    base_step.BaseStep.__init__(self, name='camera targets check', **kwargs)
    self._is_eng_build = is_eng_build
    self._product_name = product_name

  def run(self):
    eng_targets = [
      '/system/bin/camera_tool',
      '/system/bin/ffprobe',
      '/system/bin/ffmpeg',
      '/system/bin/cast_camera_hal_test'
    ]
    general_targets = [ '/obj/lib/libcast_camera_hal.so' ]
    if self._is_eng_build:
      for t in eng_targets:
        abs_path = 'out/target/product/' + self._product_name + t
        if not os.path.exists(abs_path):
          logging.error('Missing binary %s' % abs_path)
          return False
    for b in general_targets:
      abs_path = 'out/target/product/' + self._product_name + b
      if not os.path.exists(abs_path):
        logging.error('Missing binary %s' % abs_path)
        return False
    return True


class OtaUsoniaDaemonStep(OtaStep):
  """Ota Property file step class for generating repo.prop."""

  def __init__(self, build_name, **kwargs):
    """Creates a OtaPropStep instance.

    Args:
      **kwargs: Any additional args to pass to BaseStep.
    """
    OtaStep.__init__(self, build_name=build_name,
                     name='build usonia_daemon', **kwargs)

  def run(self):
    """Builds |usonia_daemon| for the |issue| and |patchset|.
    Returns:
      True if the usonia_daemon package is successfully built.
    """

    chromium_src = 'chromium/src'
    usonia_tarball = os.path.join(chromium_src,
                                  'usonia-{}.tar.gz'.format(self.build_number))
    command = ['chromecast/internal/usonia/package_wifi.py', '-o',
               '-f', '.', self.build_number]
    returncode, stdout, stderr = self.exec_subprocess(
      command, env=self.make_command_env(), cwd=chromium_src)
    if returncode != 0:
      self.exec_subprocess(['rm', '-rf', usonia_tarball])
      self._error_processor(stdout, stderr)
      return False

    archive_dir = self._ota_archive_output_dir()
    returncode, stdout, stderr = self.exec_subprocess(
      ['cp', '-r', usonia_tarball, archive_dir])

    self.exec_subprocess(['rm', '-rf', usonia_tarball])
    if returncode != 0:
      self._error_processor(stdout, stderr)
      return False

    return True


class ThreadTelemetryServiceCrosStep(OtaStep):
  """Build step class for building Thread Telemetry Service for use in Chrome
  OS."""

  def __init__(self, **kwargs):
    OtaStep.__init__(self, name='build thread_telemetry_service_cros', **kwargs)

  def run(self):
    chromium_src = 'chromium/src'
    tarball = os.path.join(chromium_src,
                           'thread_telemetry_service-{}.tar.gz'.format(
                             self.build_number))
    command = [
      'chromecast/internal/device/thread/'
      'build_thread_telemetry_service_cros.py',
      '-o',
      '-f',
      '.',
      self.build_number,
    ]
    returncode, stdout, stderr = self.exec_subprocess(
      command, env=self.make_command_env(), cwd=chromium_src)
    if returncode != 0:
      self.exec_subprocess(['rm', '-f', tarball])
      self._error_processor(stdout, stderr)
      return False

    archive_dir = self._ota_archive_output_dir()
    returncode, stdout, stderr = self.exec_subprocess(
      ['cp', '-r', tarball, archive_dir])

    self.exec_subprocess(['rm', '-f', tarball])
    if returncode != 0:
      self._error_processor(stdout, stderr)
      return False

    return True