#!/usr/bin/env vpython
# Copyright 2017 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.

"""
Takes in a page set json file. Iterate through all .wpr archives and convert
them to .wprgo archive format.

Instructions: http://bit.ly/wpr-go-migration

(1) Generate .wprgo archives from a page set json file
tools/perf/convert_legacy_wpr_archive /path/to/page_sets_json_file

(2) Test that benchmarks pass with the new page set json file.
tools/perf/run_benchmark ...

(3) Upload new .wprgo files to cloudstorage
tools/perf/convert_legacy_wpr_archive --upload /path/to/page_sets_json_file

(4) Create a CL to upload generated .wprgo.sha1 and page_sets_json_file for review

"""

import base64
import json
import optparse
import os
import re
import subprocess
import sys
import time

from core import path_util
from core import benchmark_finders

path_util.AddTelemetryToPath()
from telemetry.internal.util import binary_manager
import py_utils

path_util.AddWprToPath()
import httparchive

class JsonObject:
  def toJSON(self):
      return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4)

def dump_requests(wpr_archive):
    """List all URLs that match given params."""
    requests = '['
    requests_added = False
    for r in wpr_archive.get_requests():
        requests_added = True
        req = JsonObject()
        scheme = 'https' if r.is_ssl else 'http'
        req.url = '%s://%s%s' % (scheme, r.host, r.full_path)
        req.method = r.command
        req.headers = []
        for k in r.headers:
            h = JsonObject()
            h.key = k
            h.val = r.headers[k]
            req.headers.append(h)
        if r.request_body:
          req.body = base64.encodestring(r.request_body)
        requests += req.toJSON() + ',\n'
    if requests_added:
      requests = requests[:-2]
    requests += ']'
    return requests

def upload(page_sets_json_file, bucket):
    print "uploading to bucket %s\n" % bucket
    dir_path = os.path.abspath(os.path.dirname(page_sets_json_file))
    count = 0
    with open(page_sets_json_file, 'r+') as json_file:
      content = json_file.read()
      m = re.findall(r'"([a-zA-Z0-9_\.\-]+).wprgo"', content)
      if not m:
        raise Exception("No .wprgo files to upload in ", page_sets_json_file)
      # get rid of duplicates.
      m = list(set(m))
      for f in m:
        wprgo_file = os.path.join(dir_path, f + ".wprgo")
        wprsha1_file = os.path.join(dir_path, f + ".wpr.sha1")
        if not os.path.exists(wprgo_file):
          raise Exception('file "%s" does not exist' % wprgo_file)
        cmd = ['third_party/depot_tools/upload_to_google_storage.py', '--bucket',
               bucket, wprgo_file]
        print cmd
        return_code = subprocess.call(cmd)
        if return_code != 0:
          raise Exception("fail to upload wprgo file %s\n" % wprgo_file)
        subprocess.call(['rm', wprsha1_file])
        count += 1
    print "successfully upload %d wprgo files in %s to buckets: %s\n" % (count,
        page_sets_json_file, bucket)
    print ('Also run \"git add *.wprgo.sha1\" to include them in your CL? '
           '(Press Enter to continue or ctrl+C to skip this step)\n')
    raw_input()
    subprocess.check_call(['git', 'add', '*.wprgo.sha1'])

def GetStorySet(benchmark):
  # Create a dummy options object which hold default values that are expected
  # by Benchmark.CreateStorySet(options) method.
  parser = optparse.OptionParser()
  benchmark.AddBenchmarkCommandLineArgs(parser)
  options, _ = parser.parse_args([])
  return benchmark().CreateStorySet(options)

def main():
  option_parser = optparse.OptionParser()
  option_parser.add_option("--upload", action="store_true", default=False,
          help='Upload wprgo files to cloud storage bucket')
  options, args = option_parser.parse_args()

  if len(args) < 1:
    print 'args: %s' % args
    option_parser.error('Must specify page_sets_json_file')

  page_sets_json_file = args[0]

  if not os.path.exists(page_sets_json_file):
    option_parser.error('file "%s" does not exist' % page_sets_json_file)

  if options.upload:
    buckets = []
    benchmarks_to_skip = ['skpicture_printer_ct',
                          'screenshot_ct',
                          'repaint_ct',
                          'rasterize_and_record_micro_ct',
                          'multipage_skpicture_printer_ct',
                          'loading.cluster_telemetry',
                          'skpicture_printer',
                          'cros_tab_switching.typical_24',
                          'multipage_skpicture_printer']
    for b in benchmark_finders.GetAllBenchmarks():
      if b.Name() in benchmarks_to_skip:
        continue
      story_set = GetStorySet(b)
      if not story_set.bucket or not story_set.archive_data_file:
        continue
      if story_set.archive_data_file.split('/')[-1] == page_sets_json_file.split('/')[-1]:
        buckets.append(story_set.bucket)
    buckets = list(set(buckets))
    assert len(buckets) == 1
    upload(page_sets_json_file, buckets[0])
    return 0

  dir_path = os.path.abspath(os.path.dirname(page_sets_json_file))
  cert_file = os.path.join(path_util.GetWprGoDir(), 'wpr_cert.pem')
  key_file = os.path.join(path_util.GetWprGoDir(), 'wpr_key.pem')
  binary_manager.InitDependencyManager(None)
  go_binary_path = binary_manager.FetchPath('wpr_go',
                                            py_utils.GetHostArchName(),
                                            py_utils.GetHostOsName())
  with open(page_sets_json_file, 'r+') as json_file:
    content = json_file.read()
    m = re.findall(r'"([a-zA-Z0-9_\.\-]+).wpr"', content)
    if not m:
      raise Exception("No .wpr files found in ", page_sets_json_file)
    # get rid of duplicates.
    m = list(set(m))
    for f in m:
      wpr_file = os.path.join(dir_path, f + ".wpr")
      temp_file = os.path.join(dir_path, f + ".wprgotemp")
      wprgo_file = os.path.join(dir_path, f + ".wprgo")
      if not os.path.exists(wpr_file):
        option_parser.error('file "%s" does not exist' % wpr_file)

      archive = httparchive.HttpArchive.Load(wpr_file)
      with open(temp_file, 'w') as output:
        output.write(dump_requests(archive))
      wpr_cmd = [path_util.GetWprDir() + '/replay.py',
                 '--port=8080', '--ssl_port=8089', '--no-dns_forwarding',
                 '--inject_scripts=',
                 wpr_file]
      print wpr_cmd
      wpr_server = subprocess.Popen(' '.join(wpr_cmd), shell=True)
      time.sleep(10)
      go_cmd = [go_binary_path,
                'convert',
                '--input_file=%s' % temp_file,
                '--output_file=%s' % wprgo_file,
                '--http_port=8080',
                '--https_port=8089',
                '--https_cert_file=%s' % cert_file,
                '--https_key_file=%s' % key_file]
      return_code = subprocess.call(go_cmd)
      wpr_server.terminate()
      if return_code != 0:
        print "fail to start wpr go\n"
        return 1
      content = content.replace(f + ".wpr", f + ".wprgo")
      print "successfully written %s \n" %  f + ".wprgo"
      os.remove(temp_file)
    json_file.seek(0)
    json_file.write(content)
    json_file.truncate()
  return 0

if __name__ == '__main__':
  sys.exit(main())