#!/usr/bin/env python

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

# For Spreadsheets:
try:
  from xml.etree import ElementTree
except ImportError:
  from elementtree import ElementTree
import gdata.spreadsheet.service
import gdata.service
import atom.service
import gdata.spreadsheet
import atom

# For Issue Tracker:
import gdata.projecthosting.client
import gdata.projecthosting.data
import gdata.gauth
import gdata.client
import gdata.data
import atom.http_core
import atom.core

# For this script:
import getpass
from optparse import OptionParser
import pickle
from sets import Set

# Settings
credentials_store = 'creds.dat'

class Merger(object):
  def __init__(self, ss_key, ss_ws_key, tracker_message, tracker_project,
               debug, pretend):
    self.ss_key = ss_key
    self.ss_ws_key = ss_ws_key
    self.tracker_message = tracker_message
    self.tracker_project = tracker_project
    self.debug_enabled = debug
    self.pretend = pretend
    self.user_agent = 'adlr-tracker-spreadsheet-merger'
    self.it_keys = ['id', 'owner', 'status', 'title']

  def debug(self, message):
    """Prints message if debug mode is set."""
    if self.debug_enabled:
      print message

  def print_feed(self, feed):
    'Handy for debugging'
    for i, entry in enumerate(feed.entry):
      print 'id:', entry.id
      if isinstance(feed, gdata.spreadsheet.SpreadsheetsCellsFeed):
        print '%s %s\n' % (entry.title.text, entry.content.text)
      elif isinstance(feed, gdata.spreadsheet.SpreadsheetsListFeed):
        print '%s %s %s' % (i, entry.title.text, entry.content.text)
        # Print this row's value for each column (the custom dictionary is
        # built using the gsx: elements in the entry.)
        print 'Contents:'
        for key in entry.custom:
          print '  %s: %s' % (key, entry.custom[key].text)
        print '\n',
      else:
        print '%s %s\n' % (i, entry.title.text)


  def tracker_login(self):
    """Logs user into Tracker, using cached credentials if possible.
    Saves credentials after login."""
    self.it_client = gdata.projecthosting.client.ProjectHostingClient()
    self.it_client.source = self.user_agent

    self.load_creds()

    if self.tracker_token and self.tracker_user:
      print 'Using existing credential for tracker login'
      self.it_client.auth_token = self.tracker_token
    else:
      self.tracker_user = raw_input('Issue Tracker Login:')
      password = getpass.getpass('Password:')
      self.it_client.ClientLogin(self.tracker_user, password,
                                 source=self.user_agent, service='code',
                                 account_type='GOOGLE')
      self.tracker_token = self.it_client.auth_token
      self.store_creds()

  def spreadsheet_login(self):
    """Logs user into Google Spreadsheets, using cached credentials if possible.
    Saves credentials after login."""
    self.gd_client = gdata.spreadsheet.service.SpreadsheetsService()
    self.gd_client.source = self.user_agent

    self.load_creds()
    if self.docs_token:
      print 'Using existing credential for docs login'
      self.gd_client.SetClientLoginToken(self.docs_token)
    else:
      self.gd_client.email = raw_input('Google Docs Login:')
      self.gd_client.password = getpass.getpass('Password:')
      self.gd_client.ProgrammaticLogin()
      self.docs_token = self.gd_client.GetClientLoginToken()
      self.store_creds()

  def fetch_spreadsheet_issues(self):
    """Fetches all issues from the user-specified spreadsheet. Returns
    them as an array or dictionaries."""
    feed = self.gd_client.GetListFeed(self.ss_key, self.ss_ws_key)
    issues = []
    for entry in feed.entry:
      issue = {}
      for key in entry.custom:
        issue[key] = entry.custom[key].text
      issue['__raw_entry'] = entry
      issues.append(issue)
    return issues

  def ids_for_spreadsheet_issues(self, ss_issues):
    """Returns a Set of strings, each string an id from ss_issues"""
    ret = Set()
    for ss_issue in ss_issues:
      ret.add(ss_issue['id'])
    return ret

  def tracker_issues_for_query_feed(self, feed):
    """Converts a feed object from a query to a list of tracker issue
    dictionaries."""
    issues = []
    for issue in feed.entry:
      issue_dict = {}
      issue_dict['labels'] = [label.text for label in issue.label]
      issue_dict['id'] = issue.id.text.split('/')[-1]
      issue_dict['title'] = issue.title.text
      issue_dict['status'] = issue.status.text
      if issue.owner:
        issue_dict['owner'] = issue.owner.username.text
      issues.append(issue_dict)
    return issues

  def fetch_tracker_issues(self, ss_issues):
    """Fetches all relevant issues from traacker and returns them as an array
    of dictionaries. Relevance is:
    - has an ID that's in ss_issues, OR
    - (is Area=Installer AND status is open).
    Open status is one of: Unconfirmed, Untriaged, Available, Assigned,
    Started, Upstream"""
    issues = []
    got_results = True
    index = 1
    while got_results:
      query = gdata.projecthosting.client.Query(label='Area-Installer',
                                                max_results=50,
                                                start_index=index)
      feed = self.it_client.get_issues('chromium-os', query=query)
      if not feed.entry:
        got_results = False
      index = index + len(feed.entry)
      issues.extend(self.tracker_issues_for_query_feed(feed))
    # Now, remove issues that are open or in ss_issues.
    ss_ids = self.ids_for_spreadsheet_issues(ss_issues)
    open_statuses = ['Unconfirmed', 'Untriaged', 'Available', 'Assigned',
                     'Started', 'Upstream']
    new_issues = []
    for issue in issues:
      if issue['status'] in open_statuses or issue['id'] in ss_ids:
        new_issues.append(issue)
      # Remove id from ss_ids, if it's there
      ss_ids.discard(issue['id'])
    issues = new_issues

    # Now, for each ss_id that didn't turn up in the query, explicitly add it
    for id_ in ss_ids:
      query = gdata.projecthosting.client.Query(issue_id=id_,
                                                max_results=50,
                                                start_index=index)
      feed = self.it_client.get_issues('chromium-os', query=query)
      if not feed.entry:
        print 'No result for id', id_
        continue
      issues.extend(self.tracker_issues_for_query_feed(feed))

    return issues

  def store_creds(self):
    """Stores login credentials to disk."""
    obj = {}
    if self.docs_token:
      obj['docs_token'] = self.docs_token
    if self.tracker_token:
      obj['tracker_token'] = self.tracker_token
    if self.tracker_user:
      obj['tracker_user'] = self.tracker_user
    try:
      f = open(credentials_store, 'w')
      pickle.dump(obj, f)
      f.close()
    except IOError:
      print 'Unable to store credentials'

  def load_creds(self):
    """Loads login credentials from disk."""
    self.docs_token = None
    self.tracker_token = None
    self.tracker_user = None
    try:
      f = open(credentials_store, 'r')
      obj = pickle.load(f)
      f.close()
      if obj.has_key('docs_token'):
        self.docs_token = obj['docs_token']
      if obj.has_key('tracker_token'):
        self.tracker_token = obj['tracker_token']
      if obj.has_key('tracker_user'):
        self.tracker_user = obj['tracker_user']
    except IOError:
      print 'Unable to load credentials'

  def browse(self):
    """Browses Spreadsheets to help the user find the spreadsheet and
    worksheet keys"""
    print 'Browsing spreadsheets...'

    if self.ss_key and self.ss_ws_key:
      print 'You already passed in --ss_key and --ss_ws_key. No need to browse.'
      return

    print 'Logging in...'
    self.spreadsheet_login()

    if not self.ss_key:
      print 'Fetching spreadsheets...'
      feed = self.gd_client.GetSpreadsheetsFeed()
      print ''
      print 'Spreadsheet key - Title'
      for entry in feed.entry:
        key = entry.id.text.split('/')[-1]
        title = entry.title.text
        print '"%s" - "%s"' % (key, title)
      print ''
      print 'Done. Rerun with --ss_key=KEY to browse a list of worksheet keys.'
    else:
      print 'Fetching worksheets for spreadsheet', self.ss_key
      feed = self.gd_client.GetWorksheetsFeed(self.ss_key)
      for entry in feed.entry:
        key = entry.id.text.split('/')[-1]
        title = entry.title.text
        print ''
        print 'Worksheet key - Title'
        print '"%s" - "%s"' % (key, title)
      print ''
      print 'Done. You now have keys for --ss_key and --ss_ws_key.'

  def tracker_issue_for_id(self, issues, id_):
    """Returns the element of issues which has id_ for the key 'id'"""
    for issue in issues:
      if issue['id'] == id_:
        return issue
    return None

  def spreadsheet_issue_to_tracker_dict(self, ss_issue):
    """Converts a spreadsheet issue to the dict format that is used to
    represent a tracker issue."""
    ret = {}
    ret['project'] = self.tracker_project
    ret['title'] = ss_issue['title']
    ret['summary'] = self.tracker_message
    ret['owner'] = ss_issue['owner']
    if ss_issue.get('status') is not None:
      ret['status'] = ss_issue['status']
    ret['labels'] = []
    for (key, value) in ss_issue.items():
      if key.endswith('-') and (value is not None):
        ret['labels'].append(key.title() + value)
    return ret

  def label_from_prefix(self, prefix, corpus):
    """Given a corpus (array of lable strings), return the first label
    that begins with the specified prefix."""
    for label in corpus:
      if label.startswith(prefix):
        return label
    return None

  def update_spreadsheet_issue_to_tracker_dict(self, ss_issue, t_issue):
    """Updates a given tracker issue with data from the spreadsheet issue."""
    ret = {}
    ret['title'] = ss_issue['title']
    ret['id'] = ss_issue['id']
    ret['summary'] = self.tracker_message
    if ss_issue['status'] != t_issue['status']:
      ret['status'] = ss_issue['status']

    if ss_issue.get('owner'):
      if (not t_issue.has_key('owner')) or \
        (ss_issue['owner'] != t_issue['owner']):
        ret['owner'] = ss_issue['owner']
    # labels
    ret['labels'] = []
    for (key, value) in ss_issue.items():
      caps_key = key.title()
      if not caps_key.endswith('-'):
        continue
      ss_label = None
      if value:
        ss_label = caps_key + value.title()
      t_label = self.label_from_prefix(caps_key, t_issue['labels'])

      if t_label is None and ss_label is None:
        # Nothing
        continue

      if (t_label is not None) and \
        ((ss_label is None) or (ss_label != t_label)):
        ret['labels'].append('-' + t_label)

      if (ss_label is not None) and \
        ((t_label is None) or (t_label != ss_label)):
        ret['labels'].append(ss_label)
    return ret

  def tracker_issue_has_changed(self, t_issue, ss_issue):
    """Returns True iff ss_issue indicates changes in t_issue that need to be
    committed up to the Issue Tracker."""
    if t_issue is None:
      return True
    potential_commit = \
      self.update_spreadsheet_issue_to_tracker_dict(ss_issue, t_issue)

    if potential_commit.has_key('status') or \
      potential_commit.has_key('owner') or \
      (len(potential_commit['labels']) > 0):
      return True
    if potential_commit['title'] != t_issue['title']:
      return True
    return False

  def spreadsheet_to_tracker_commits(self, ss_issues, t_issues):
    """Given the current state of all spreadsheet issues and tracker issues,
    returns a list of all commits that need to go to tracker to get it in
    line with the spreadsheet."""
    ret = []
    for ss_issue in ss_issues:
      t_issue = self.tracker_issue_for_id(t_issues, ss_issue['id'])
      commit = {}
      # TODO see if an update is needed at all
      if t_issue is None:
        commit['type'] = 'append'
        commit['dict'] = self.spreadsheet_issue_to_tracker_dict(ss_issue)
        commit['__ss_issue'] = ss_issue
      else:
        if not self.tracker_issue_has_changed(t_issue, ss_issue):
          continue
        commit['type'] = 'update'
        commit['dict'] = \
          self.update_spreadsheet_issue_to_tracker_dict(ss_issue, t_issue)
      ret.append(commit)
    return ret

  def fetch_issues(self):
    """Logs into Docs/Tracker, and fetches spreadsheet and tracker issues"""
    print 'Logging into Docs...'
    self.spreadsheet_login()
    print 'Logging into Tracker...'
    self.tracker_login()

    print 'Fetching spreadsheet issues...'
    ss_issues = self.fetch_spreadsheet_issues()
    self.debug('Spreadsheet issues: %s' % ss_issues)
    print 'Fetching tracker issues...'
    t_issues = self.fetch_tracker_issues(ss_issues)
    self.debug('Tracker issues: %s' % t_issues)
    return (t_issues, ss_issues)

  def spreadsheet_to_tracker(self):
    """High-level function to manage migrating data from the spreadsheet
    to Tracker."""
    (t_issues, ss_issues) = self.fetch_issues()
    print 'Calculating deltas...'
    commits = self.spreadsheet_to_tracker_commits(ss_issues, t_issues)
    self.debug('got commits: %s' % commits)
    if not commits:
      print 'No deltas. Done.'
      return

    for commit in commits:
      dic = commit['dict']
      labels = dic.get('labels')
      owner = dic.get('owner')
      status = dic.get('status')

      if commit['type'] == 'append':
        print 'Creating new tracker issue...'
        if self.pretend:
          print '(Skipping because --pretend is set)'
          continue
        created = self.it_client.add_issue(self.tracker_project,
                                           dic['title'],
                                           self.tracker_message,
                                           self.tracker_user,
                                           labels=labels,
                                           owner=owner,
                                           status=status)
        issue_id = created.id.text.split('/')[-1]
        print 'Created issue with id:', issue_id
        print 'Write id back to spreadsheet row...'
        raw_entry = commit['__ss_issue']['__raw_entry']
        ss_issue = commit['__ss_issue']
        del ss_issue['__raw_entry']
        ss_issue.update({'id': issue_id})
        self.gd_client.UpdateRow(raw_entry, ss_issue)
        print 'Done.'
      else:
        print 'Updating issue with id:', dic['id']
        if self.pretend:
          print '(Skipping because --pretend is set)'
          continue
        self.it_client.update_issue(self.tracker_project,
                                    dic['id'],
                                    self.tracker_user,
                                    comment=self.tracker_message,
                                    status=status,
                                    owner=owner,
                                    labels=labels)
        print 'Done.'

  def spreadsheet_issue_for_id(self, issues, id_):
    """Given the array of spreadsheet issues, return the first one that
    has id_ for the key 'id'."""
    for issue in issues:
      if issue['id'] == id_:
        return issue
    return None

  def value_for_key_in_labels(self, label_array, prefix):
    """Given an array of labels and a prefix, return the non-prefix part
    of the first label that has that prefix. E.g. if label_array is
    ["Mstone-R7", "Area-Installer"] and prefix is "Area-", returns
    "Installer"."""
    for label in label_array:
      if label.startswith(prefix):
        return label[len(prefix):]
    return None

  def tracker_issue_to_spreadsheet_issue(self, t_issue, ss_keys):
    """Converts a tracker issue to the format used by spreadsheet, given
    the row headings ss_keys."""
    new_row = {}
    for key in ss_keys:
      if key.endswith('-'):
        # label
        new_row[key] = self.value_for_key_in_labels(t_issue['labels'],
                                                    key.title())
      # Special cases
      if key in self.it_keys and key in t_issue:
        new_row[key] = t_issue[key]
    return new_row

  def spreadsheet_row_needs_update(self, ss_issue, t_issue):
    """Returns True iff the spreadsheet issue passed in needs to be updated
    to match data in the tracker issue."""
    new_ss_issue = self.tracker_issue_to_spreadsheet_issue(t_issue,
                                                           ss_issue.keys())
    for key in new_ss_issue.keys():
      if not ss_issue.has_key(key):
        continue
      if new_ss_issue[key] != ss_issue[key]:
        return True
    return False

  def tracker_to_spreadsheet_commits(self, t_issues, ss_issues):
    """Given the current set of spreadsheet and tracker issues, computes
    commits needed to go to Spreadsheets to get the spreadsheet in line
    with what's in Tracker."""
    ret = []
    keys = ss_issues[0].keys()
    for t_issue in t_issues:
      commit = {}
      ss_issue = self.spreadsheet_issue_for_id(ss_issues, t_issue['id'])
      if ss_issue is None:
        # New issue
        commit['new_row'] = self.tracker_issue_to_spreadsheet_issue(t_issue,
                                                                    keys)
        commit['type'] = 'append'
      elif self.spreadsheet_row_needs_update(ss_issue, t_issue):
        commit['__raw_entry'] = ss_issue['__raw_entry']
        del ss_issue['__raw_entry']
        ss_issue.update(self.tracker_issue_to_spreadsheet_issue(t_issue, keys))
        commit['dict'] = ss_issue
        commit['type'] = 'update'
      else:
        continue
      ret.append(commit)
    return ret

  def tracker_to_spreadsheet(self):
    """High-level function to migrate data from Tracker to the spreadsheet."""
    (t_issues, ss_issues) = self.fetch_issues()
    if len(ss_issues) == 0:
      raise Exception('Error: must have at least one non-header row in '\
                      'spreadsheet')
      return
    ss_keys = ss_issues[0].keys()

    print 'Calculating deltas...'
    ss_commits = self.tracker_to_spreadsheet_commits(t_issues, ss_issues)
    self.debug('commits: %s' % ss_commits)
    if not ss_commits:
      print 'Nothing to commit.'
      return
    print 'Committing...'
    for commit in ss_commits:
      self.debug('Operating on commit: %s' % commit)
      if commit['type'] == 'append':
        print 'Appending new row...'
        if not self.pretend:
          self.gd_client.InsertRow(commit['new_row'],
                                   self.ss_key, self.ss_ws_key)
        else:
          print '(Skipped because --pretend set)'
      if commit['type'] == 'update':
        print 'Updating row...'
        if not self.pretend:
          self.gd_client.UpdateRow(commit['__raw_entry'], commit['dict'])
        else:
          print '(Skipped because --pretend set)'
      print 'Done.'

def main():
  class PureEpilogOptionParser(OptionParser):
    def format_epilog(self, formatter):
      return self.epilog

  parser = PureEpilogOptionParser()
  parser.add_option('-a', '--action', dest='action', metavar='ACTION',
                    help='Action to perform')
  parser.add_option('-d', '--debug', action='store_true', dest='debug',
                    default=False, help='Print debug output.')
  parser.add_option('-m', '--message', dest='message', metavar='TEXT',
                    help='Log message when updating Tracker issues')
  parser.add_option('-p', '--pretend', action='store_true', dest='pretend',
                    default=False, help="Don't commit anything.")
  parser.add_option('--ss_key', dest='ss_key', metavar='KEY',
                    help='Spreadsheets key (find with browse action)')
  parser.add_option('--ss_ws_key', dest='ss_ws_key', metavar='KEY',
                    help='Spreadsheets worksheet key (find with browse action)')
  parser.add_option('--tracker_project', dest='tracker_project',
                    metavar='PROJECT',
                    help='Tracker project (default: chromium-os)',
                    default='chromium-os')
  parser.epilog = """Actions:
  browse -- browse spreadsheets to find spreadsheet and worksheet keys.
  ss_to_t -- for each entry in spreadsheet, apply its values to tracker.
             If no ID is in the spreadsheet row, a new tracker item is created
             and the spreadsheet is updated.
  t_to_ss -- for each tracker entry, apply it or add it to the spreadsheet.


This script can be used to migrate Issue Tracker issues between Issue Tracker
and Google Spreadsheets. The spreadsheet should have certain columns in any
order: Id, Owner, Title, Status. The spreadsheet may have any label of the
form 'Key-'. For those labels that end in '-', this script assumes the cell
value and the header form a label that should be applied to the issue. E.g.
if the spredsheet has a column named 'Mstone-' and a cell under it called
'R8' that corresponds to the label 'Mstone-R8' in Issue Tracker.

To migrate data, you must choose on each invocation of this script if you
wish to migrate data from Issue Tracker to a spreadsheet of vice-versa.

When migrating from Tracker, all found issues based on the query
(which is currently hard-coded to "label=Area-Installer") will be inserted
into the spreadsheet (overwritng existing cells if a row with matching ID
is found). Custom columns in the spreadsheet won't be overwritten, so if
the spreadsheet contains extra columns about issues (e.g. time estimates)
they will be preserved.

When migrating from spreadsheet to Tracker, each row in the spreadsheet
is compared to existing tracker issues that match the query
(which is currently hard-coded to "label=Area-Installer"). If the
spreadsheet row has no Id, a new Issue Tracker issue is created and the new
Id is written back to the spreadsheet. If an existing tracker issue exists,
it's updated with the data from the spreadsheet if anything has changed.

Suggested usage:
- Create a spreadsheet with columns Id, Owner, Title, Status, and any label
prefixes as desired.
- Run this script with '-b' to browse your spreadsheet and get the
spreadsheet key.
- Run this script again with '-b' and the spreadsheet key to get the
worksheet key.
- Run this script with "-a t_to_ss" or "-a ss_to_t" to migrate data in either
direction.

Known issues:
- query is currently hardcoded to label=Area-Installer. That should be
a command-line flag.
- When creating a new issue on tracker, the owner field isn't set. I (adlr)
am not sure why. Workaround: If you rerun this script, tho, it will detect
a delta and update the tracker issue with the owner, which seems to succeed.
"""

  (options, args) = parser.parse_args()

  merger = Merger(options.ss_key, options.ss_ws_key,
                  options.message, options.tracker_project,
                  options.debug, options.pretend)
  if options.action == 'browse':
    merger.browse()
  elif options.action == 'ss_to_t':
    if not options.message:
      print 'Error: when updating tracker, -m MESSAGE required.'
      return
    merger.spreadsheet_to_tracker()
  elif options.action == 't_to_ss':
    merger.tracker_to_spreadsheet()
  else:
    raise Exception('Unknown action requested.')

if __name__ == '__main__':
  main()