#!/usr/bin/env python

# Copyright (c) 2011 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Android system-wide tracing utility.

This is a tool for capturing a trace that includes data from both userland and
the kernel.  It creates an HTML file for visualizing the trace.
"""

import errno, optparse, os, select, subprocess, sys, time, zlib

# This list is based on the tags in frameworks/native/include/utils/Trace.h.
trace_tag_bits = {
  'gfx':      1<<1,
  'input':    1<<2,
  'view':     1<<3,
  'webview':  1<<4,
  'wm':       1<<5,
  'am':       1<<6,
  'sync':     1<<7,
  'audio':    1<<8,
  'video':    1<<9,
  'camera':   1<<10,
}

flattened_html_file = 'systrace_trace_viewer.html'

def add_adb_serial(command, serial):
  if serial != None:
    command.insert(1, serial)
    command.insert(1, '-s')

def main():
  parser = optparse.OptionParser()
  parser.add_option('-o', dest='output_file', help='write HTML to FILE',
                    default='trace.html', metavar='FILE')
  parser.add_option('-t', '--time', dest='trace_time', type='int',
                    help='trace for N seconds', metavar='N')
  parser.add_option('-b', '--buf-size', dest='trace_buf_size', type='int',
                    help='use a trace buffer size of N KB', metavar='N')
  parser.add_option('-d', '--disk', dest='trace_disk', default=False,
                    action='store_true', help='trace disk I/O (requires root)')
  parser.add_option('-f', '--cpu-freq', dest='trace_cpu_freq', default=False,
                    action='store_true', help='trace CPU frequency changes')
  parser.add_option('-i', '--cpu-idle', dest='trace_cpu_idle', default=False,
                    action='store_true', help='trace CPU idle events')
  parser.add_option('-l', '--cpu-load', dest='trace_cpu_load', default=False,
                    action='store_true', help='trace CPU load')
  parser.add_option('-s', '--no-cpu-sched', dest='trace_cpu_sched', default=True,
                    action='store_false', help='inhibit tracing CPU ' +
                    'scheduler (allows longer trace times by reducing data ' +
                    'rate into buffer)')
  parser.add_option('-u', '--bus-utilization', dest='trace_bus_utilization',
                    default=False, action='store_true',
                    help='trace bus utilization (requires root)')
  parser.add_option('-w', '--workqueue', dest='trace_workqueue', default=False,
                    action='store_true', help='trace the kernel workqueues ' +
                    '(requires root)')
  parser.add_option('--set-tags', dest='set_tags', action='store',
                    help='set the enabled trace tags and exit; set to a ' +
                    'comma separated list of: ' +
                    ', '.join(trace_tag_bits.iterkeys()))
  parser.add_option('--link-assets', dest='link_assets', default=False,
                    action='store_true', help='link to original CSS or JS resources '
                    'instead of embedding them')
  parser.add_option('--from-file', dest='from_file', action='store',
                    help='read the trace from a file rather than running a live trace')
  parser.add_option('--asset-dir', dest='asset_dir', default='trace-viewer',
                    type='string', help='')
  parser.add_option('-e', '--serial', dest='device_serial', type='string',
                    help='adb device serial number')
  options, args = parser.parse_args()

  if options.link_assets or options.asset_dir != 'trace-viewer':
    parser.error('--link-assets and --asset-dir is deprecated.')

  if options.set_tags:
    flags = 0
    tags = options.set_tags.split(',')
    for tag in tags:
      try:
        flags |= trace_tag_bits[tag]
      except KeyError:
        parser.error('unrecognized tag: %s\nknown tags are: %s' %
                     (tag, ', '.join(trace_tag_bits.iterkeys())))
    atrace_args = ['adb', 'shell', 'setprop', 'debug.atrace.tags.enableflags', hex(flags)]
    add_adb_serial(atrace_args, options.device_serial)
    try:
      subprocess.check_call(atrace_args)
    except subprocess.CalledProcessError, e:
      print >> sys.stderr, 'unable to set tags: %s' % e
    print '\nSet enabled tags to: %s\n' % ', '.join(tags)
    print ('You will likely need to restart the Android framework for this to ' +
          'take effect:\n\n    adb shell stop\n    adb shell ' +
          'start\n')
    return

  atrace_args = ['adb', 'shell', 'atrace', '-z']
  add_adb_serial(atrace_args, options.device_serial)

  if options.trace_disk:
    atrace_args.append('-d')
  if options.trace_cpu_freq:
    atrace_args.append('-f')
  if options.trace_cpu_idle:
    atrace_args.append('-i')
  if options.trace_cpu_load:
    atrace_args.append('-l')
  if options.trace_cpu_sched:
    atrace_args.append('-s')
  if options.trace_bus_utilization:
    atrace_args.append('-u')
  if options.trace_workqueue:
    atrace_args.append('-w')
  if options.trace_time is not None:
    if options.trace_time > 0:
      atrace_args.extend(['-t', str(options.trace_time)])
    else:
      parser.error('the trace time must be a positive number')
  if options.trace_buf_size is not None:
    if options.trace_buf_size > 0:
      atrace_args.extend(['-b', str(options.trace_buf_size)])
    else:
      parser.error('the trace buffer size must be a positive number')

  if options.from_file is not None:
    atrace_args = ['cat', options.from_file]

  script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))

  with open(os.path.join(script_dir, flattened_html_file), 'r') as f:
    trace_viewer_html = f.read()

  html_filename = options.output_file

  trace_started = False
  leftovers = ''
  adb = subprocess.Popen(atrace_args, stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
  dec = zlib.decompressobj()
  while True:
    ready = select.select([adb.stdout, adb.stderr], [], [adb.stdout, adb.stderr])
    if adb.stderr in ready[0]:
      err = os.read(adb.stderr.fileno(), 4096)
      sys.stderr.write(err)
      sys.stderr.flush()
    if adb.stdout in ready[0]:
      out = leftovers + os.read(adb.stdout.fileno(), 4096)
      if options.from_file is None:
        out = out.replace('\r\n', '\n')
      if out.endswith('\r'):
        out = out[:-1]
        leftovers = '\r'
      else:
        leftovers = ''
      if not trace_started:
        lines = out.splitlines(True)
        out = ''
        for i, line in enumerate(lines):
          if line == 'TRACE:\n':
            sys.stdout.write("downloading trace...")
            sys.stdout.flush()
            out = ''.join(lines[i+1:])
            html_prefix = read_asset(script_dir, 'prefix.html')
            html_file = open(html_filename, 'w')
            html_file.write(
              html_prefix.replace("{{SYSTRACE_TRACE_VIEWER_HTML}}",
                  trace_viewer_html))
            html_file.write('<!-- BEGIN TRACE -->\n' +
              '  <script class="trace-data" type="application/text">\n')
            trace_started = True
            break
          elif 'TRACE:'.startswith(line) and i == len(lines) - 1:
            leftovers = line + leftovers
          else:
            sys.stdout.write(line)
            sys.stdout.flush()
      if len(out) > 0:
        out = dec.decompress(out)
      html_out = out.replace('\n', '\\n\\\n')
      if len(html_out) > 0:
        html_file.write(html_out)
    result = adb.poll()
    if result is not None:
      break
  if result != 0:
    print >> sys.stderr, 'adb returned error code %d' % result
  elif trace_started:
    html_out = dec.flush().replace('\n', '\\n\\\n').replace('\r', '')
    if len(html_out) > 0:
      html_file.write(html_out)
    html_file.write('  </script>\n<!-- END TRACE -->\n')
    html_suffix = read_asset(script_dir, 'suffix.html')
    html_file.write(html_suffix)
    html_file.close()
    print " done\n\n    wrote file://%s\n" % (os.path.abspath(options.output_file))
  else:
    print >> sys.stderr, ('An error occured while capturing the trace.  Output ' +
      'file was not written.')

def read_asset(src_dir, filename):
  return open(os.path.join(src_dir, filename)).read()

if __name__ == '__main__':
  main()