#!/usr/bin/env python
# Copyright 2018 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This script parses result of build/android/adb_profile_chrome_startup and
# prints various information.

from __future__ import print_function

import argparse
import collections
import glob
import itertools
import json
import os
import re


def HumanSortingKey(string):
  # From https://goo.gl/65xrva
  def _ToInt(string):
    return int(string) if string.isdigit() else string

  return [_ToInt(c) for c in re.split('(\d+)', string)]


class LogIndenter(object):
  _indentation = 0

  def __init__(self, message=None, *arguments):
    self._indented = False
    if message is not None:
      log(message, *arguments)

  def __enter__(self):
    self.indent()
    return self

  def __exit__(self, type, value, traceback):
    self.unindent()

  def indent(self):
    if not self._indented:
      LogIndenter._indentation += 1
      self._indented = True

  def unindent(self):
    if self._indented:
      LogIndenter._indentation -= 1
      self._indented = False

  @classmethod
  def indentation(self):
    return LogIndenter._indentation


def log(message, *arguments):
  if not message:
    print()
    return

  if arguments:
    message = message.format(*arguments)
  if LogIndenter.indentation() > 0:
    message = '  ' * LogIndenter.indentation() + message
  print(message)


def ParseTraceDatas(trace_file_path):
  start_tag = re.compile(
    '^\s*<script class="trace-data" type="application/text">$'
  )
  end_tag = re.compile('^(?P<line>.*?)\s*</script>$')

  trace_datas = []

  current_trace_lines = None
  with open(trace_file_path) as trace_file:
    for line in trace_file:
      line = line.rstrip()
      if current_trace_lines is None:
        if start_tag.match(line):
          current_trace_lines = []
      else:
        match = end_tag.match(line)
        if match:
          current_trace_lines.append(match.group('line'))
          trace_datas.append('\n'.join(current_trace_lines))
          current_trace_lines = None
        else:
          current_trace_lines.append(line)

  return trace_datas


class Event(object):
  PHASE_BEGIN = 'B'
  PHASE_END = 'E'
  PHASE_COMPLETE = 'X'
  PHASE_ASYNC_BEGIN = 'S'
  PHASE_ASYNC_END = 'F'

  def __init__(self, node):
    self._node = node

  @property
  def pid(self):
    return int(self._node['pid'])

  @property
  def tid(self):
    return int(self._node['tid'])

  @property
  def name(self):
    return self._node.get('name')

  @property
  def phase(self):
    return self._node['ph']

  @property
  def category(self):
    return self._node['cat']

  @property
  def timestamp_us(self):
    return long(self._node['ts'])

  @property
  def duration_us(self):
    return long(self._node['dur'])

  @property
  def args(self):
    return self._node['args']


class EventInterval(object):
  def __init__(self, from_event=None, to_event=None):
    self.from_event = from_event
    self.to_event = to_event

  def SetFromEventOnce(self, from_event):
    if self.from_event is None:
      self.from_event = from_event

  def SetToEventOnce(self, to_event):
    if self.to_event is None:
      self.to_event = to_event

  def FormatAsMilliseconds(self):
    if not self.from_event:
      time_string = "[missing the start of of the interval]"
    elif not self.to_event:
      time_string = "[missing the end of the interval]"
    else:
      interval_us = self.to_event.timestamp_us - self.from_event.timestamp_us
      time_string = str(interval_us / 1000.0)
    return time_string

  def UpTo(self, other_iterval):
    return EventInterval(self.from_event, other_iterval.to_event)


class Process(object):
  BROWSER_NAME = 'Browser'

  def __init__(self, pid):
    self.pid = pid
    self.name = None
    self.events_by_name = collections.defaultdict(list)
    self.time_ns_by_histogram = {}
    self.malloc_counter_by_name = {}
    # TODO: move these into Trace
    self.startup_interval = EventInterval()
    self.first_ui_interval = EventInterval()


class Trace(object):
  def __init__(self, file_path):
    self.file_path = file_path
    self.process_by_pid = {}
    self.startup_event = None
    self.navigation_start_event = None
    # TODO: convert these to properties over events
    self.navigation_to_contentul_paint_interval = EventInterval()
    self.navigation_to_meaningful_paint_interval = EventInterval()
    self.navigation_to_commit_interval = None

  @property
  def startup_to_navigation_interval(self):
    return EventInterval(self.startup_event, self.navigation_start_event)

  def Finalize(self):
    self.startup_event = self.FindFirstEvent(*Trace.STARTUP_EVENT_NAMES)
    self.navigation_start_event = self.FindFirstEvent(
      Trace.NAVIGATION_START_EVENT_NAME
    )

    def _FindNavigationToCommitInterval():
      events = self.FindAllEvents(Trace.NAVIGATION_COMMIT_EVENT_NAME)
      interval = EventInterval()
      for event in events:
        if event.phase == Event.PHASE_ASYNC_BEGIN:
          interval.SetFromEventOnce(event)
        elif event.phase == Event.PHASE_ASYNC_END:
          interval.SetToEventOnce(event)
      return interval

    self.navigation_to_commit_interval = _FindNavigationToCommitInterval()

  def FindAllEvents(self, *names):
    events = []
    for process in self.process_by_pid.itervalues():
      for name in names:
        process_events = process.events_by_name.get(name)
        if process_events:
          events.extend(process_events)
    events.sort(key=lambda e: e.timestamp_us)
    return events

  def FindFirstEvent(self, *names):
    events = self.FindAllEvents(*names)
    return events[0] if events else None

  NAVIGATION_START_EVENT_NAME = 'NavigationTiming navigationStart'
  NAVIGATION_COMMIT_EVENT_NAME = 'Navigation StartToCommit'

  STARTUP_EVENT_NAMES = [
    'Startup.BrowserMainEntryPoint',
    'ChromeApplication.onCreate',
    'ContentShellApplication.onCreate',
  ]


def ParseTrace(file_path):
  trace_datas = ParseTraceDatas(file_path)
  if not trace_datas:
    raise Exception("The file doesn't have any trace-data elements.")

  trace_json = None
  for trace_data in trace_datas:
    try:
      trace_json = json.loads(trace_data)
    except ValueError:
      continue

  if not trace_json:
    raise Exception("Couldn't parse trace-data json.")

  trace = Trace(file_path)

  for event_node in trace_json['traceEvents']:
    event = Event(event_node)

    pid = event.pid
    process = trace.process_by_pid.get(event.pid)
    if not process:
      process = Process(pid)
      trace.process_by_pid[pid] = process

    name = event.name
    if not name:
      continue

    process.events_by_name[name].append(event)

    phase = event.phase
    category = event.category

    if name == 'process_name':
      process.name = event.args['name']

    if (
      category == 'disabled-by-default-uma-addtime'
      and name not in process.time_ns_by_histogram
    ):
      process.time_ns_by_histogram[name] = int(event.args['value_ns'])

    if name in Trace.STARTUP_EVENT_NAMES:
      process.startup_interval.SetFromEventOnce(event)
    elif name == 'BenchmarkInstrumentation::ImplThreadRenderingStats':
      process.startup_interval.SetToEventOnce(event)

    if name == Trace.NAVIGATION_START_EVENT_NAME:
      trace.navigation_to_contentul_paint_interval.SetFromEventOnce(event)
      trace.navigation_to_meaningful_paint_interval.SetFromEventOnce(event)
    elif name == 'firstContentfulPaint':
      trace.navigation_to_contentul_paint_interval.SetToEventOnce(event)
    elif name == 'firstMeaningfulPaint':
      trace.navigation_to_meaningful_paint_interval.SetToEventOnce(event)

    if (
      name == 'AsyncInitializationActivity.onCreate()'
      and phase == Event.PHASE_END
    ):
      process.first_ui_interval.SetFromEventOnce(event)
    elif name == 'ChromeBrowserInitializer.startChromeBrowserProcessesAsync':
      process.first_ui_interval.SetToEventOnce(event)

    if category == 'malloc' and name == 'malloc_counter':
      counter_name, counter_value = next(event.args.iteritems())
      process.malloc_counter_by_name[counter_name] = long(counter_value)

  trace.Finalize()
  return trace


EventSummary = collections.namedtuple(
  'EventSummary',
  [
    'trace',
    'event',
    'startup_to_event_ms',
    'navigation_to_event_ms',
    'duration_ms',
  ],
)


def SummarizeEvents(event_name_regex, trace, process):
  summaries = []

  def _AddSummary(event, start_us, duration_us):
    startup_to_event_ms = (
      None
      if trace.startup_event is None
      else (start_us - trace.startup_event.timestamp_us) / 1000.0
    )
    navigation_to_event_ms = (
      None
      if trace.navigation_start_event is None
      else (start_us - trace.navigation_start_event.timestamp_us) / 1000.0
    )
    summaries.append(
      EventSummary(
        trace,
        event,
        startup_to_event_ms,
        navigation_to_event_ms,
        duration_us / 1000.0,
      )
    )

  for name, events in process.events_by_name.iteritems():
    if event_name_regex.search(name):
      sorted_events = sorted(events, key=lambda e: (e.tid, e.timestamp_us))
      begin_event = None
      for event in sorted_events:
        if event.phase == Event.PHASE_COMPLETE:
          _AddSummary(event, event.timestamp_us, event.duration_us)
        elif (
          event.phase == Event.PHASE_BEGIN
          or event.phase == Event.PHASE_ASYNC_BEGIN
        ):
          begin_event = event
        elif (
          event.phase == Event.PHASE_END or event.phase == Event.PHASE_ASYNC_END
        ):
          if begin_event is not None:
            duration_us = event.timestamp_us - begin_event.timestamp_us
            _AddSummary(event, begin_event.timestamp_us, duration_us)
          begin_event = None

  return summaries


def PrintReport(file_paths, options):
  # TODO: don't accumulate traces, build report on the fly
  traces = []
  for file_path in file_paths:
    log('Parsing {}...', file_path)
    try:
      traces.append(ParseTrace(file_path))
    except Exception as e:
      log('Oops: {}', e.message)

  log('Parsed {} trace(s).', len(traces))

  event_name_regex = None
  event_summaries_by_name = collections.defaultdict(list)
  if options.print_events:
    event_name_regex = re.compile(options.print_events)

  def _TraceSortingKey(trace):
    return HumanSortingKey(os.path.basename(trace.file_path))

  traces.sort(key=lambda t: _TraceSortingKey(t))

  if options.csv:
    separator = ','
    gap = ''
  else:
    separator = '\t'
    # Make it less likely for terminals to eat tabs when wrapping a line.
    gap = '    '

  table = [
    [
      'File',
      'Startup (ms)',
      'StartupToNavigation (ms)',
      'NavigationToCommit (ms)',
      'NavigationToContentfulPaint (ms)',
      'StartupToContentfulPaint (ms)',
      'NavigationToMeaningfulPaint (ms)',
      'StartupToMeaningfulPaint (ms)',
    ]
  ]
  for trace in traces:
    browser_process = None
    for process in trace.process_by_pid.itervalues():
      if process.name == Process.BROWSER_NAME:
        browser_process = process
        break
    if browser_process is None:
      continue

    table.append(
      [
        os.path.basename(trace.file_path),
        browser_process.startup_interval.FormatAsMilliseconds(),
        trace.startup_to_navigation_interval.FormatAsMilliseconds(),
        trace.navigation_to_commit_interval.FormatAsMilliseconds(),
        trace.navigation_to_contentul_paint_interval.FormatAsMilliseconds(),
        browser_process.startup_interval.UpTo(
          trace.navigation_to_contentul_paint_interval
        ).FormatAsMilliseconds(),
        trace.navigation_to_meaningful_paint_interval.FormatAsMilliseconds(),
        browser_process.startup_interval.UpTo(
          trace.navigation_to_meaningful_paint_interval
        ).FormatAsMilliseconds(),
      ]
    )

    if event_name_regex:
      event_summaries = SummarizeEvents(
        event_name_regex, trace, browser_process
      )
      for summary in event_summaries:
        event_summaries_by_name[summary.event.name].append(summary)

  for name, event_summaries in event_summaries_by_name.iteritems():
    table.append([])

    summaries_by_trace = collections.defaultdict(list)
    for summary in event_summaries:
      summaries_by_trace[summary.trace].append(summary)

    width = max(len(s) for s in summaries_by_trace.itervalues())
    summary_headers = [
      'StartupToEvent (ms)',
      'NavigationToEvent (ms)',
      'Duration (ms)',
    ]

    table.append(
      [name]
      + ([gap] * len(summary_headers))
      + list(
        itertools.chain.from_iterable(
          ['#{}'.format(i)] + [gap] * (len(summary_headers) - 1)
          for i in range(1, width)
        )
      )
    )
    table.append(['File'] + summary_headers * width)

    trace_summaries = sorted(
      summaries_by_trace.iteritems(), key=lambda t_s: _TraceSortingKey(t_s[0])
    )
    for trace, summaries in trace_summaries:
      row = [os.path.basename(trace.file_path)]
      for summary in summaries:
        row += [
          (
            gap
            if summary.startup_to_event_ms is None
            else summary.startup_to_event_ms
          ),
          (
            gap
            if summary.navigation_to_event_ms is None
            else summary.navigation_to_event_ms
          ),
          summary.duration_ms,
        ]
      table.append(row)

  print()
  print('\n'.join(separator.join(str(v) for v in row) for row in table))


def PrintTrace(trace_file_path, options):
  trace = ParseTrace(trace_file_path)

  def _PrintInterval(name, interval):
    log('{} (ms): {}', name, interval.FormatAsMilliseconds())

  def _PrintHistogramTime(process, name):
    time_ns = process.time_ns_by_histogram.get(name)
    time_ms = None if time_ns is None else time_ns / 1e6
    if time_ms is not None or options.print_none_histograms:
      log('{} (ms): {}', name, time_ms)

  histogram_names = [
    'ChromeGeneratedCustomTab.IntentToFirstCommitNavigationTime3.ZoomedIn',
    'CustomTabs.IntentToFirstCommitNavigationTime3.ZoomedIn',
    'PageLoad.PaintTiming.NavigationToFirstPaint',
    'PageLoad.PaintTiming.NavigationToFirstContentfulPaint',
    'PageLoad.Experimental.PaintTiming.NavigationToFirstMeaningfulPaint',
    'SessionRestore.ForegroundTabFirstPaint3',
  ]

  processes = sorted(trace.process_by_pid.itervalues(), key=lambda p: p.name)

  events_regex = None
  if options.print_events:
    events_regex = re.compile(options.print_events)

  for process in processes:
    log('{} timings:', process.name)
    indenter = LogIndenter()
    indenter.indent()

    if process.name == Process.BROWSER_NAME:
      _PrintInterval('Startup', process.startup_interval)
      _PrintInterval(
        'StartupToNavigation', trace.startup_to_navigation_interval
      )
      _PrintInterval('NavigationToCommit', trace.navigation_to_commit_interval)
      _PrintInterval(
        'NavigationToContentfulPaint',
        trace.navigation_to_contentul_paint_interval,
      )
      _PrintInterval(
        'StartupToContentfulPaint',
        process.startup_interval.UpTo(
          trace.navigation_to_contentul_paint_interval
        ),
      )
      _PrintInterval(
        'NavigationToMeaningfulPaint',
        trace.navigation_to_meaningful_paint_interval,
      )
      _PrintInterval(
        'StartupToMeaningfulPaint',
        process.startup_interval.UpTo(
          trace.navigation_to_meaningful_paint_interval
        ),
      )

      if options.experimental:
        _PrintInterval('First UI interval', process.first_ui_interval)

    if process.malloc_counter_by_name:

      def _PrintMallocCounter(title, value_name, factor):
        value = process.malloc_counter_by_name.get(value_name)
        if value is not None:
          value /= factor
        log('{}: {}', title, value)

      log('Malloc counters:')
      with LogIndenter():
        _PrintMallocCounter('Total time (ms)', 'malloc_time_ns', 1000000)
        _PrintMallocCounter('Total allocated (KiB)', 'allocated_bytes', 1024)
        _PrintMallocCounter('Allocations', 'allocation_count', 1)
        _PrintMallocCounter('Frees', 'free_count', 1)

    for histogram_name in histogram_names:
      _PrintHistogramTime(process, histogram_name)

    if events_regex:
      event_summaries = SummarizeEvents(events_regex, trace, process)
      if event_summaries:
        with LogIndenter('Events matching "{}":', events_regex.pattern):
          for event_summary in event_summaries:
            with LogIndenter('{}:', event_summary.event.name):
              if event_summary.startup_to_event_ms is not None:
                log(
                  'StartupToEvent (ms): {}', event_summary.startup_to_event_ms
                )
              if event_summary.navigation_to_event_ms is not None:
                log(
                  'NavigationToEvent (ms): {}',
                  event_summary.navigation_to_event_ms,
                )
              log('Duration (ms): {}', event_summary.duration_ms)

    indenter.unindent()
    log('')


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('file_or_glob')
  parser.add_argument(
    '--print-none-histograms',
    help='Print histograms with None values.',
    default=False,
    action='store_true',
  )
  # TODO: introduce a variant that takes a list of event names, as escaping
  #       event names can be tedious.
  # TODO: match regex against '<index>|<event name>' to allow selecting
  #       events by index (complicated for begin/end pairs).
  parser.add_argument(
    '--print-events', help='Print events matching the specified regex.'
  )
  parser.add_argument(
    '--experimental',
    default=False,
    action='store_true',
    help='Enable experimental stuff.',
  )
  parser.add_argument(
    '--report',
    default=False,
    action='store_true',
    help='Present information as a tab-separated table.',
  )
  parser.add_argument(
    '--csv',
    default=False,
    action='store_true',
    help=('Separate report values by commas (not tabs).'),
  )

  options = parser.parse_args()

  globbed = False
  if os.path.isfile(options.file_or_glob):
    trace_file_paths = [options.file_or_glob]
  else:
    globbed = True
    file_pattern = options.file_or_glob
    trace_file_paths = glob.glob(file_pattern)
    if not trace_file_paths:
      file_pattern += '*html'
      trace_file_paths = glob.glob(file_pattern)
    if not trace_file_paths:
      log("'{}' didn't match anything.", file_pattern)
      return
    log("'{}' matched {} file(s).", file_pattern, len(trace_file_paths))
    log('')

  if options.report:
    PrintReport(trace_file_paths, options)
  else:
    for file_path in trace_file_paths:
      if globbed:
        log('_' * len(file_path))
        log(file_path)
      log('')
      PrintTrace(file_path, options)


if __name__ == '__main__':
  main()