# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Utilities for code coverage related processings."""

import logging
import os
import posixpath
import shutil
import subprocess

from devil import base_error
from pylib import constants

# These are use for code coverage.
LLVM_PROFDATA_PATH = os.path.join(constants.DIR_SOURCE_ROOT, 'third_party',
                                  'llvm-build', 'Release+Asserts', 'bin',
                                  'llvm-profdata')
# Name of the file extension for profraw data files.
_PROFRAW_FILE_EXTENSION = 'profraw'
# Name of the file where profraw data files are merged.
_MERGE_PROFDATA_FILE_NAME = 'coverage_merged.' + _PROFRAW_FILE_EXTENSION


def GetDeviceClangCoverageDir(device):
  """Gets the directory to generate clang coverage data on device.

  Args:
    device: The working device.

  Returns:
    The directory path on the device.
  """
  return posixpath.join(device.GetExternalStoragePath(), 'chrome', 'test',
                        'coverage', 'profraw')


def PullAndMaybeMergeClangCoverageFiles(device, device_coverage_dir, output_dir,
                                        output_subfolder_name):
  """Pulls and possibly merges clang coverage file to a single file.

  Only merges when llvm-profdata tool exists. If so, Merged file is at
  `output_dir/coverage_merged.profraw`and raw profraw files before merging
  are deleted.

  Args:
    device: The working device.
    device_coverage_dir: The directory storing coverage data on device.
    output_dir: The output directory on host to store the
        coverage_merged.profraw file.
    output_subfolder_name: The subfolder in |output_dir| to pull
        |device_coverage_dir| into. It will be deleted after merging if
        merging happens.
  """
  if not device.PathExists(device_coverage_dir, retries=0):
    logging.warning('Clang coverage data folder does not exist on device: %s',
                    device_coverage_dir)
    return
  # Host side dir to pull device coverage profraw folder into.
  profraw_parent_dir = os.path.join(output_dir, output_subfolder_name)
  # Note: The function pulls |device_coverage_dir| folder,
  # instead of profraw files, into |profraw_parent_dir|. the
  # function also removes |device_coverage_dir| from device.
  PullClangCoverageFiles(device, device_coverage_dir, profraw_parent_dir)
  # Merge data into one merged file if llvm-profdata tool exists.
  if os.path.isfile(LLVM_PROFDATA_PATH):
    profraw_folder_name = os.path.basename(
        os.path.normpath(device_coverage_dir))
    profraw_dir = os.path.join(profraw_parent_dir, profraw_folder_name)
    MergeClangCoverageFiles(output_dir, profraw_dir)
    shutil.rmtree(profraw_parent_dir)


def PullClangCoverageFiles(device, device_coverage_dir, output_dir):
  """Pulls clang coverage files on device to host directory.

  Args:
    device: The working device.
    device_coverage_dir: The directory to store coverage data on device.
    output_dir: The output directory on host.
  """
  try:
    if not os.path.exists(output_dir):
      os.makedirs(output_dir)
    device.PullFile(device_coverage_dir, output_dir)
    if not os.listdir(os.path.join(output_dir, 'profraw')):
      logging.warning('No clang coverage data was generated for this run')
  except (OSError, base_error.BaseError) as e:
    logging.warning('Failed to pull clang coverage data, error: %s', e)
  finally:
    device.RemovePath(device_coverage_dir, force=True, recursive=True)


def MergeClangCoverageFiles(coverage_dir, profdata_dir):
  """Merge coverage data files.

  Each instrumentation activity generates a separate profraw data file. This
  merges all profraw files in profdata_dir into a single file in
  coverage_dir. This happens after each test, rather than waiting until after
  all tests are ran to reduce the memory footprint used by all the profraw
  files.

  Args:
    coverage_dir: The path to the coverage directory.
    profdata_dir: The directory where the profraw data file(s) are located.

  Return:
    None
  """
  # profdata_dir may not exist if pulling coverage files failed.
  if not os.path.exists(profdata_dir):
    logging.debug('Profraw directory does not exist: %s', profdata_dir)
    return

  merge_file = os.path.join(coverage_dir, _MERGE_PROFDATA_FILE_NAME)
  profraw_files = [
      os.path.join(profdata_dir, f) for f in os.listdir(profdata_dir)
      if f.endswith(_PROFRAW_FILE_EXTENSION)
  ]

  try:
    logging.debug('Merging target profraw files into merged profraw file.')
    subprocess_cmd = [
        LLVM_PROFDATA_PATH,
        'merge',
        '-o',
        merge_file,
        '-sparse=true',
    ]
    # Grow the merge file by merging it with itself and the new files.
    if os.path.exists(merge_file):
      subprocess_cmd.append(merge_file)
    subprocess_cmd.extend(profraw_files)
    output = subprocess.check_output(subprocess_cmd).decode('utf8')
    logging.debug('Merge output: %s', output)

  except subprocess.CalledProcessError:
    # Don't raise error as that will kill the test run. When code coverage
    # generates a report, that will raise the error in the report generation.
    logging.error(
        'Failed to merge target profdata files to create '
        'merged profraw file for files: %s', profraw_files)

  # Free up memory space on bot as all data is in the merge file.
  for f in profraw_files:
    os.remove(f)