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

"""Creates a Python telemetry page_set from the specified webpages CSV.

This module does the following steps:
* Downloads a ZIP from http://s3.amazonaws.com/alexa-static/top-1m.csv.zip
* Unpacks it and reads its contents in memory.
* Writes out multiple Python page sets from the CSV file for the specified
  number of webpages.

Sample Usage:
  python create_page_set.py -s 1

Running the above command will create a page set with the webpage in 1st
position in the CSV.
"""

__author__ = 'Ravi Mistry'

import json
import optparse
import os
import urllib
import zipfile

from StringIO import StringIO


if '__main__' == __name__:
  option_parser = optparse.OptionParser()
  option_parser.add_option(
      '-s', '--position',
      help='Specifies the start position of the webpages in the CSV which will '
           'be created as pagesets.',
      default='1')
  option_parser.add_option(
      '-e', '--end',
      help='Specifies the end position of the webpages in the CSV which will '
           'be created as pagesets.',
      default='1')
  option_parser.add_option(
      '-c', '--csv_file',
      help='Location of a filtered alexa top 1M CSV file. Each row should '
           'have 3 entries, 1st will be rank, 2nd will be domain name and '
           'third will be the fully qualified url. If the third section is '
           'missing then a page_set for the URL will not be generated. If '
           'csv_file is not specified then this script downloads it from the '
           'internet.',
      default=None)
  option_parser.add_option(
      '-p', '--pagesets_type',
      help='The type of pagesets to create from the 1M list. Eg: All, '
           '100k, 10k, IndexSample10k, Mobile10k',
      default='All')
  option_parser.add_option(
      '-u', '--useragent_type',
      help='The type of user agent to use in the pagesets. Eg: desktop, '
           'mobile, tablet',
      default='desktop')
  option_parser.add_option(
      '-o', '--pagesets_output_dir',
      help='The directory generated pagesets will be outputted in.',
      default='page_sets')
  options, unused_args = option_parser.parse_args()

  # Validate arguments.
  if int(options.position) <= 0:
    raise Exception('The -s/--position must be greater than 0')
  if int(options.end) < int(options.position):
    raise Exception('The -e/--end must be >= than -s/--position')

  start_index = int(options.position) - 1
  with open(options.csv_file) as fp:
    for i, line in enumerate(fp):
      if i < start_index:
        continue
      elif i >= int(options.end):
        break
      else:
        websites = []
        website = line.strip().split(',')[1]
        if 'PDF' in options.pagesets_type:
          # PDF urls do not need any additional prefixes.
          qualified_website = website
        elif website.startswith('https://') or website.startswith('http://'):
          qualified_website = website
        else:
          qualified_website = 'http://www.%s' % website
        websites.append(qualified_website)

        archive_data_file = os.path.join(
            '/', 'b', 'storage', 'webpage_archives',
            options.pagesets_type,
            '%s.json' % options.position)

        page_set_content = {
          'user_agent': options.useragent_type,
          'archive_data_file': archive_data_file,
          'urls_list': ','.join(websites),
        }

        # Output the pageset to a file.
        parent_dir = os.path.join(options.pagesets_output_dir, str(i+1))
        os.mkdir(parent_dir)
        with open(os.path.join(parent_dir, '%s.py' % (i+1)), 'w') as outfile:
          json.dump(page_set_content, outfile)

@@ -271,251 +271,72 @@
 ule 
-does the following steps:%0A* Downloads a ZIP from http://s3.amazonaws.com/alexa-static/top-1m.csv.zip%0A* Unpacks it and reads its contents in memory.%0A* Writes out multiple Python page sets from the CSV file for the specified%0A  number of webpages
+writes out multiple Python page sets from the specified CSV file
 .%0A%0AS
@@ -380,16 +380,39 @@
 .py -s 1
+ -e 10 -c /tmp/test.csv
 %0A%0ARunnin
@@ -443,17 +443,18 @@
  create 
-a
+10
  page se
@@ -458,33 +458,51 @@
  set
- with the webpage in 1st%0A
+s containing webpages from the%0A1st to 10th 
 posi
@@ -1067,32 +1067,51 @@
 ted as pagesets.
+ This is inclusive.
 ',%0A      default

Improve documentation of create_page_set.py