#!/usr/bin/python2.6
# Copyright (c) 2011 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.

"""Support uploading a csv file to a Google Docs spreadsheet."""

import optparse
import os
import pickle
import re
import sys

import gdata.spreadsheet.service

sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
import chromite.lib.table as table
import chromite.lib.cros_build_lib as cros_lib
import chromite.lib.upgrade_table as utable
import merge_package_status as mps

REAL_SS_KEY='tJuuSuHmrEMqdL5b8dkgBIA'
TEST_SS_KEY='t3RE08XLO2f1vTiV4N-w2ng'

def _PrepColForSS(col):
  """Translate a column name for spreadsheet "list" interface."""
  # Spreadsheet "list" interface requires column names to be
  # all lowercase and with no spaces.
  return col.lower().replace(' ', '')

def _PrepRowForSS(row):
  """Make sure spreadsheet handles all values in row as strings."""
  for key in row:
    val = row[key]
    row[key] = _PrepValForSS(val)
  return row

# Regex to detect values that the spreadsheet will auto-format as numbers.
_NUM_REGEX = re.compile(r'^[\d\.]+$')
def _PrepValForSS(val):
  """Make sure spreadsheet handles this value as a string."""
  if val and _NUM_REGEX.match(val):
    return "'" + val
  return val

def _ScrubValFromSS(val):
  """Remove string indicator prefix if found."""
  if val and val[0] == "'":
    return val[1:]
  return val

class Uploader(object):
  """Uploads portage package status data from csv file to Google spreadsheet."""

  __slots__ = ['_docs_token', # Google Docs login token
               '_email',      # Google Docs account email
               '_gd_client',  # gdata.spreadsheet.service.SpreadsheetsService
               '_password',   # Google Docs account password
               '_ss_key',     # Google Doc spreadsheet key (string)
               '_table',      # table.Table
               '_verbose',    # boolean
               '_ws_key',     # Worksheet key (string)
               ]

  ID_COL = utable.UpgradeTable.COL_PACKAGE
  SOURCE = "Uploaded from CSV"

  def __init__(self, table, verbose=False):
    self._table = table
    self._gd_client = None
    self._docs_token = None
    self._email = None
    self._password = None
    self._ss_key = None
    self._verbose = verbose
    self._ws_key = None

  def _Verbose(self, msg):
    """Print |msg| if _verbose is true."""
    if self._verbose:
      print msg

  def _GetWorksheetKey(self, ss_key, ws_index):
    """Get the worksheet key at index |ws_index| in spreadsheet |ss_key|."""
    feed = self._gd_client.GetWorksheetsFeed(ss_key)
    # The worksheet key is the last component in the URL (after last '/')
    return feed.entry[ws_index].id.text.split('/')[-1]

  def LoginDocsWithEmailPassword(self):
    """Set up and connect the Google Doc client using email/password."""
    gd_client = gdata.spreadsheet.service.SpreadsheetsService()
    gd_client.source = self.SOURCE
    gd_client.email = self._email
    gd_client.password = self._password
    gd_client.ProgrammaticLogin()
    self._gd_client = gd_client
    self._docs_token = gd_client.GetClientLoginToken()

  def LoginDocsWithToken(self):
    """Set up and connect Google Doc client using login token."""
    print "Using existing credentials for docs login."
    gd_client = gdata.spreadsheet.service.SpreadsheetsService()
    gd_client.source = self.SOURCE
    gd_client.SetClientLoginToken(self._docs_token)
    self._gd_client = gd_client
    self._email = None
    self._password = None

  def SetCreds(self, email, password):
    """Set the email/password credentials to be used."""
    self._email = email
    self._password = password

  def LoadCreds(self, filepath):
    """Store email/password credentials to |filepath|."""
    print "Loading Docs login credentials from '%s'" % filepath
    # Read email from first line and password from second.
    f = open(filepath, 'r')
    (self._email, self._password) = (l.strip() for l in f.readlines())
    f.close()

  def StoreCreds(self, filepath):
    """Load email/password credentials from |filepath|."""
    print "Storing Docs login credentials to '%s'" % filepath
    # Simply write email on first line and password on second.
    f = open(filepath, 'w')
    f.write(self._email + '\n')
    f.write(self._password + '\n')
    f.close()

  def LoadDocsToken(self, filepath):
    """Loads docs login token from |filepath|."""
    self._docs_token = None
    try:
      print "Loading Docs login token from '%s'" % filepath
      f = open(filepath, 'r')
      obj = pickle.load(f)
      f.close()
      if obj.has_key('docs_token'):
        self._docs_token = obj['docs_token']
    except IOError:
      print 'Unable to load credentials from file at %s' % filepath

    return bool(self._docs_token)

  def StoreDocsToken(self, filepath):
    """Stores docs login token to disk at |filepath|."""
    obj = {}
    if self._docs_token:
      obj['docs_token'] = self._docs_token
    try:
      print "Storing Docs login token to '%s'" % filepath
      f = open(filepath, 'w')
      pickle.dump(obj, f)
      f.close()
    except IOError:
      print 'Unable to store credentials'

  def _GetSSRowForPackage(self, package):
    """Find the spreadsheet row object corresponding to Package=|package|."""
    query = gdata.spreadsheet.service.ListQuery()
    query._SetSpreadsheetQuery('package = "%s"' % package)
    feed = self._gd_client.GetListFeed(self._ss_key, self._ws_key, query=query)
    if len(feed.entry) == 1:
      return feed.entry[0]

    if len(feed.entry) > 1:
      raise LookupError("More than one row in spreadsheet with Package=%s" %
                        package)

    if len(feed.entry) == 0:
      return None

  def Upload(self, ss_key, ws_index=0, ws_key=None):
    """Upload |_table| to the given Google Spreadsheet.

    The spreadsheet is identified the spreadsheet key |ss_key|.
    The worksheet within that spreadsheet is identified by either a worksheet
    key |ws_key| if given, or a worksheet index |ws_index|.  An index of 0,
    for example, means the first worksheet in that spreadsheet.
    """
    self._ss_key = ss_key

    if ws_key is None:
      ws_key = self._GetWorksheetKey(ss_key, ws_index)
    self._ws_key = ws_key

    self._Verbose("Details by package: S=Same, C=Changed, A=Added, D=Deleted")

    rows_unchanged, rows_updated, rows_inserted = self._UploadChangedRows()

    rows_deleted, rows_with_owner_deleted = self._DeleteOldRows()

    print("Final row stats: %d changed, %d added, %d deleted, %d same." %
          (rows_updated, rows_inserted, rows_deleted, rows_unchanged))
    if rows_with_owner_deleted:
      print("WARNING: %d rows with owner entry deleted, see above warnings." %
            rows_with_owner_deleted)
    else:
      print("No rows with owner entry were deleted.")

  def _UploadChangedRows(self):
    """Upload all rows in table that need to be changed in spreadsheet."""
    print "Uploading to spreadsheet '%s' now." % self._ss_key

    rows_unchanged, rows_updated, rows_inserted = (0, 0, 0)

    # Go over all rows in csv table.  Identify existing row by the 'Package'
    # column.  Either update existing row or create new one.
    for csv_row in self._table:
      csv_package = csv_row[self.ID_COL]
      ss_row = self._GetSSRowForPackage(csv_package)

      # Seed new row values from csv_row values, with column translation.
      new_row = dict((_PrepColForSS(key), csv_row[key]) for key in csv_row)

      if ss_row:
        changed = []

        # For columns that are in spreadsheet but not in csv, grab values
        # from spreadsheet so they are not overwritten by UpdateRow.
        for col in ss_row.custom:
          ss_val = _ScrubValFromSS(ss_row.custom[col].text)
          if col not in new_row:
            new_row[col] = ss_val
          elif ss_val and new_row[col] and ss_val != new_row[col]:
            changed.append("%s='%s'->'%s'" % (col, ss_val, new_row[col]))
          elif ss_row.custom[col].text != _PrepValForSS(ss_val):
            changed.append("%s=str'%s'" % (col, new_row[col]))

        if changed:
          self._gd_client.UpdateRow(ss_row, _PrepRowForSS(new_row))
          rows_updated += 1
          self._Verbose("C %-30s: %s" %
                        (csv_package, ', '.join(changed)))
        else:
          rows_unchanged += 1
          self._Verbose("S %-30s:" % csv_package)
      else:
        self._gd_client.InsertRow(_PrepRowForSS(new_row),
                                  self._ss_key, self._ws_key)
        rows_inserted += 1
        row_descr_list = []
        for col in sorted(new_row.keys()):
          if col != self.ID_COL:
            row_descr_list.append("%s='%s'" % (col, new_row[col]))
        self._Verbose("A %-30s: %s" %
                      (csv_package, ', '.join(row_descr_list)))

    return (rows_unchanged, rows_updated, rows_inserted)

  def _DeleteOldRows(self):
    """Delete all rows from spreadsheet that not found in table."""
    print "Checking for rows in spreadsheet that should be deleted now."

    rows_deleted, rows_with_owner_deleted = (0, 0)

    # Also need to delete rows in spreadsheet that are not in csv table.
    list_feed = self._gd_client.GetListFeed(self._ss_key, self._ws_key)
    for ss_row in list_feed.entry:
      ss_id_col = _PrepColForSS(self.ID_COL)
      ss_package = _ScrubValFromSS(ss_row.custom[ss_id_col].text)

      # See whether this row is in csv table.
      csv_rows = self._table.GetRowsByValue({ self.ID_COL: ss_package })
      if not csv_rows:
        # Row needs to be deleted from spreadsheet.
        owner_val = None
        owner_notes_val = None
        row_descr_list = []
        for col in sorted(ss_row.custom.keys()):
          if col == 'owner':
            owner_val = ss_row.custom[col].text
          if col == 'ownernotes':
            owner_notes_val = ss_row.custom[col].text

          # Don't include ID_COL value in description, it is in prefix already.
          if col != _PrepColForSS(self.ID_COL):
            val = ss_row.custom[col].text
            row_descr_list.append("%s='%s'" % (col, val))

        self._Verbose("D %-30s: %s" %
                      (ss_package, ', '.join(row_descr_list)))
        if owner_val or owner_notes_val:
          rows_with_owner_deleted += 1
          print("WARNING: Deleting spreadsheet row with owner entry:\n" +
                "  %-30s: Owner=%s, Owner Notes=%s" %
                (ss_package, owner_val, owner_notes_val))

        self._gd_client.DeleteRow(ss_row)
        rows_deleted += 1

    return (rows_deleted, rows_with_owner_deleted)

def LoadTable(table_file):
  """Load csv |table_file| into a table.  Return table."""
  print "Loading csv table from '%s'." % (table_file)
  csv_table = table.Table.LoadFromCSV(table_file)
  return csv_table

def main():
  """Main function."""
  usage = 'Usage: %prog [options] csv_file'
  parser = optparse.OptionParser(usage=usage)
  parser.add_option('--cred-file', dest='cred_file', type='string',
                    action='store', default=None,
                    help="File for reading/writing Docs login email/password.")
  parser.add_option('--email', dest='email', type='string',
                    action='store', default=None,
                    help="Email for Google Doc user")
  parser.add_option('--password', dest='password', type='string',
                    action='store', default=None,
                    help="Password for Google Doc user")
  parser.add_option('--ss-key', dest='ss_key', type='string',
                    action='store', default=None,
                    help="Key of spreadsheet to upload to")
  parser.add_option('--test-spreadsheet', dest='test_ss',
                    action='store_true', default=False,
                    help="Upload to the testing spreadsheet.")
  parser.add_option('--token-file', dest='token_file', type='string',
                    action='store', default=None,
                    help="File for reading/writing Docs login token.")
  parser.add_option('--verbose', dest='verbose',
                    action='store_true', default=False,
                    help="Show details about packages..")

  (options, args) = parser.parse_args()

  if len(args) < 1:
    parser.print_help()
    cros_lib.Die("One csv_file is required.")

  # If email or password provided, the other is required.  If neither is
  # provided, then either cred_file or token_file must be provided and
  # be a real file.
  if options.email or options.password:
    if not (options.email and options.password):
      parser.print_help()
      cros_lib.Die("The email/password options must be used together.")
  elif not options.token_file and not options.cred_file:
    parser.print_help()
    cros_lib.Die("Either email/password, cred-file, or token-file required.")
  elif not ((options.token_file and os.path.exists(options.token_file)) or
            (options.cred_file and os.path.exists(options.cred_file))):
    parser.print_help()
    cros_lib.Die("Without email/password cred-file or token-file must exist.")
  elif (options.token_file and os.path.exists(options.token_file) and
        options.cred_file and os.path.exists(options.cred_file)):
    parser.print_help()
    cros_lib.Die("Without email/password, both cred-file and token-file " +
                 "is ambiguous.")

  # --ss-key and --test-spreadsheet are mutually exclusive.
  if options.ss_key and options.test_ss:
    parser.print_help()
    cros_lib.Die("Cannot specify --ss-key and --test-spreadsheet together.")

  # Load the given csv file.
  csv_table = LoadTable(args[0])

  # Prepare table for upload.
  mps.FinalizeTable(csv_table)

  # Prepare the Google Doc client for uploading.
  uploader = Uploader(csv_table, verbose=options.verbose)
  if options.email and options.password:
    # Login with supplied email/password
    uploader.SetCreds(email=options.email, password=options.password)
    uploader.LoginDocsWithEmailPassword()

    # If cred or token file given, write them out now.
    if options.token_file:
      uploader.StoreDocsToken(options.token_file)
    if options.cred_file:
      uploader.StoreCreds(options.cred_file)
  elif options.cred_file and os.path.exists(options.cred_file):
    # Login by reading email/password from supplied file.
    uploader.LoadCreds(options.cred_file)
    uploader.LoginDocsWithEmailPassword()

    # If token file given, write it out now.
    if options.token_file:
      uploader.StoreDocsToken(options.token_file)
  elif options.token_file and os.path.exists(options.token_file):
    # Login with token from supplied file.
    uploader.LoadDocsToken(options.token_file)
    uploader.LoginDocsWithToken()

  ss_key = options.ss_key
  if not ss_key:
    if options.test_ss:
      ss_key = TEST_SS_KEY # For testing with backup spreadsheet
    else:
      ss_key = REAL_SS_KEY
  uploader.Upload(ss_key)

if __name__ == '__main__':
  main()