#!/usr/bin/env python
# Copyright (c) 2014 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.

"""Module that downloads SKPs from Google Storage into a specified dir.

Downloads SKPs which were generated by Cluster Telemetry's GCE slaves using:
http://skia-tree-status.appspot.com/skia-telemetry
"""

import optparse
import os
import posixpath
import sys

# This number is unlikely to ever change.
# TODO(rmistry): Read NUMBER_OF_GCE_SLAVES and DEST_GSBASE from a Cluster
# Telemetry global_variables.json
NUMBER_OF_GCE_SLAVES = 100
# This bucket is also unlikely to ever change (though it would be great if it
# did).
DEST_GSBASE = 'gs://chromium-skia-gm'

PARENT_DIR = os.path.dirname(os.path.realpath(__file__))
CHROMIUM_SITE_CONFIG_DIR = os.path.join(PARENT_DIR, os.pardir, 'third_party',
                                        'chromium_buildbot', 'site_config')

# Set the PYTHONPATH for this script to include chromium_buildbot scripts
# and site_config. Also include skia buildbot's slave utils directory.
sys.path.append(
    os.path.join(PARENT_DIR, os.pardir, 'third_party', 'chromium_buildbot',
                 'scripts'))
sys.path.append(CHROMIUM_SITE_CONFIG_DIR)
sys.path.append(
    os.path.join(PARENT_DIR, os.pardir, 'slave', 'skia_slave_scripts', 'utils'))

import gs_utils


def DownloadSKPs(chromium_build, pageset_type, dest_dir):

  # Validate that the .boto file exists.
  if not os.path.isfile(os.path.join(CHROMIUM_SITE_CONFIG_DIR, '.boto')):
    raise Exception(
        'Please copy your .boto file to %s.' % CHROMIUM_SITE_CONFIG_DIR)

  gs_skp_slave_template = posixpath.join(
      DEST_GSBASE, 'telemetry', 'skps', 'slave%s', pageset_type, chromium_build)

  # Validate that the requested SKP directory exists for worker1.
  if not gs_utils.does_storage_object_exist(gs_skp_slave_template % '1'):
    raise Exception(
        'There are no SKPs generated for the pageset: %s and the '
        'chromium_build: %s' % (pageset_type, chromium_build))

  # Create destination directory if it does not exist.
  if not os.path.isdir(dest_dir):
    os.makedirs(dest_dir)

  # Ensure that the specified destination directory is empty.
  if os.listdir(dest_dir):
    raise Exception('%s is not empty. Please provide an empty dir.' % dest_dir)

  # Transfer one slave's generated SKPs at a time to the specified directory.
  for slave_num in range(1, NUMBER_OF_GCE_SLAVES + 1):
    # Create a directory for the slave to download to.
    local_slave_dir = os.path.join(dest_dir, chromium_build,
                                   'slave%s' % slave_num)
    os.makedirs(local_slave_dir)
    # Download from Google Storage to the local slave directory.
    gs_utils.copy_storage_directory(
        src_dir=posixpath.join(gs_skp_slave_template % slave_num, '*'),
        dest_dir=local_slave_dir)


if '__main__' == __name__:
  option_parser = optparse.OptionParser()
  option_parser.add_option(
    '-b', '--chromium_build',
    help='Specifies which chromium build was used to generate the SKPs. '
         'Eg: e69f2ce-e2b1569. The list of all builds are at: '
         'http://skia-tree-status.appspot.com/skia-telemetry/chromium_builds')
  option_parser.add_option(
    '-p', '--pageset_type',
    help='Specifies the pageset type we want to download. Options are '
         'All/100k/10k/IndexSample10k. Defaults to All.',
    default='All')
  option_parser.add_option(
    '-d', '--dest_dir',
    help='Specifies the directory where the SKPs will be downloaded too. %s '
         'slave directories will be created each containing their own '
         'corresponding SKPs.' % NUMBER_OF_GCE_SLAVES)
  options, unused_args = option_parser.parse_args()

  # Validate required parameters.
  assert options.chromium_build, 'Must specify --chromium_build'
  assert options.dest_dir, 'Must specify --dest_dir'

  DownloadSKPs(options.chromium_build, options.pageset_type, options.dest_dir)