# Copyright 2021 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Symbolizes a perfetto trace file."""

import logging
import os
import shutil
import subprocess
import sys
import tempfile

sys.path.insert(
    0,
    os.path.join(os.path.dirname(__file__), os.pardir, os.pardir, 'third_party',
                 'catapult', 'systrace'))

import breakpad_file_extractor
import flag_utils
import metadata_extractor
import rename_breakpad
import symbol_fetcher


def SymbolizeTrace(trace_file, options):
  """Symbolizes a perfetto trace file.

  Args:
    trace_file: path to proto trace file to symbolize.
    options: The options set by the command line args.

  Raises:
    Exception: if breakpad_output_dir is not empty or if path to local breakpad
      directory is invalid.
  """
  print('Symbolizing file')

  need_cleanup = False
  if options.local_breakpad_dir is None:
    breakpad_file_extractor.EnsureDumpSymsBinary(options.dump_syms_path,
                                                 options.local_build_dir)

    # Ensure valid |breakpad_output_dir|
    if options.breakpad_output_dir is None:
      # Create a temp dir if output dir is not provided.
      # Temp dir must be cleaned up later.
      options.breakpad_output_dir = tempfile.mkdtemp()
      need_cleanup = True
      flag_utils.GetTracingLogger().warning(
          'Created temporary directory to hold symbol files. These symbols '
          'will be cleaned up after symbolization. Please specify '
          '--breakpad_output_dir=<cached_dir> to save the symbols, if you need '
          'to profile multiple times. The future runs need to use '
          '--local_breakpad_dir=<cached_dir> flag so the symbolizer uses the '
          'cache.')
    else:
      if os.path.isdir(options.breakpad_output_dir):
        if os.listdir(options.breakpad_output_dir):
          raise Exception('Breakpad output directory is not empty: ' +
                          options.breakpad_output_dir)
      else:
        os.makedirs(options.breakpad_output_dir)
        flag_utils.GetTracingLogger().debug(
            'Created directory to hold symbol files.')
  else:
    if not os.path.isdir(options.local_breakpad_dir):
      raise Exception('Local breakpad directory is not valid.')
    options.breakpad_output_dir = options.local_breakpad_dir

  _EnsureBreakpadSymbols(trace_file, options)

  # Set output file to write trace data and symbols to.
  if options.output_file is None:
    options.output_file = os.path.join(
        os.path.dirname(trace_file),
        os.path.basename(trace_file) + '_symbolized_trace')

  _Symbolize(trace_file, options.symbolizer_path, options.breakpad_output_dir,
             options.output_file)

  print('Symbolized trace saved to: ' + os.path.abspath(options.output_file))

  # Cleanup
  if need_cleanup:
    flag_utils.GetTracingLogger().debug('Cleaning up symbol files.')
    shutil.rmtree(options.breakpad_output_dir)


def _EnsureBreakpadSymbols(trace_file, options):
  """Ensures that there are breakpad symbols to symbolize with.

  Args:
    trace_file: The trace file to be symbolized.
    options: The options set by the command line args. This is used to check if
      symbols need to be fetched, extracted, or if they are already present.

  Raises:
    Exception: if no breakpad files could be extracted.
  """
  # If |options.local_breakpad_dir| is not None, then this can be skipped and
  # |trace_file| can be symbolized using those symbols.
  if options.local_breakpad_dir is not None:
    return

  # Extract Metadata
  flag_utils.GetTracingLogger().info('Extracting proto trace metadata.')
  trace_metadata = metadata_extractor.MetadataExtractor(
      options.trace_processor_path, trace_file)
  trace_metadata.Initialize()
  flag_utils.GetTracingLogger().info(trace_metadata)

  if options.local_build_dir is not None:
    # Extract breakpad symbol files from binaries in |options.local_build_dir|.
    if not breakpad_file_extractor.ExtractBreakpadFiles(
        options.dump_syms_path,
        options.local_build_dir,
        options.breakpad_output_dir,
        search_unstripped=True,
        module_ids=breakpad_file_extractor.GetModuleIdsToSymbolize(
            trace_metadata)):
      raise Exception(
          'No breakpad symbols could be extracted from files in: %s xor %s' %
          (options.local_build_dir,
           os.path.join(options.local_build_dir, 'lib.unstripped')))

    rename_breakpad.RenameBreakpadFiles(options.breakpad_output_dir,
                                        options.breakpad_output_dir)
    return

  # Fetch trace breakpad symbols from GCS
  flag_utils.GetTracingLogger().info(
      'Fetching and extracting trace breakpad symbols.')
  symbol_fetcher.GetTraceBreakpadSymbols(options.cloud_storage_bucket,
                                         trace_metadata,
                                         options.breakpad_output_dir,
                                         options.dump_syms_path)


def _Symbolize(trace_file, symbolizer_path, breakpad_output_dir, output_file):
  """"Symbolizes a trace.

  Args:
    trace_file: The trace file to be symbolized.
    symbolizer_path: The path to the trace_to_text tool to use.
    breakpad_output_dir: Contains the breakpad symbols to use for symbolization.
    output_file: The path to the file to output symbols to.

  Raises:
    RuntimeError: If _RunSymbolizer() fails to execute the given command.
  """
  # Set environment variable as location of stored breakpad files.
  symbolize_env = os.environ.copy()
  symbolize_env['BREAKPAD_SYMBOL_DIR'] = os.path.join(breakpad_output_dir, '')
  cmd = [symbolizer_path, 'symbolize', trace_file]

  # Open temporary file where symbols can be stored.
  with tempfile.TemporaryFile(mode='wb+') as temp_symbol_file:
    _RunSymbolizer(cmd, symbolize_env, temp_symbol_file)

    # Write trace data and symbol data to the same file.
    temp_symbol_file.seek(0)
    symbol_data = temp_symbol_file.read()
    with open(trace_file, 'rb') as f:
      trace_data = f.read()
    with open(output_file, 'wb') as f:
      f.write(trace_data)
      f.write(symbol_data)
      flag_utils.GetTracingLogger().info(
          'Symbolized %s(%d bytes) with %d bytes of symbol data',
          os.path.abspath(trace_file), len(trace_data), len(symbol_data))


def _RunSymbolizer(cmd, env, stdout):
  proc = subprocess.Popen(cmd, env=env, stdout=stdout, stderr=subprocess.PIPE)
  out, stderr = proc.communicate()
  flag_utils.GetTracingLogger().debug('STDOUT:%s', str(out))
  flag_utils.GetTracingLogger().debug('STDERR:%s', str(stderr))
  if proc.returncode != 0:
    raise RuntimeError(str(stderr))