# Copyright 2016 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.

"""Export entities to gcloud datastore."""

from __future__ import print_function

from gcloud import datastore
import ast
import json

from chromite.lib import commandline
from chromite.lib import iter_utils

_BATCH_CHUNK_SIZE = 500


def GetParser():
  """Creates the argparse parser."""
  parser = commandline.ArgumentParser(description=__doc__)
  parser.add_argument('service_acct_json', type=str, action='store',
                      help='Path to service account credentials JSON file.')
  parser.add_argument('entities', type=str, action='store',
                      help=('Path to file with entities to export. '
                            'File should be newline-separated JSON entries.'))
  parser.add_argument('--project_id', '-i', type=str, action='store',
                      default=None,
                      help=('Optional project_id of datastore to write to. If '
                            'not supplied, will be taken from credentials '
                            'file.'))
  parser.add_argument('--namespace', '-n', type=str, action='store',
                      default=None,
                      help='Optional namespace in which to store entities.')
  parser.add_argument('--parent_key', '-p', type=str, action='store',
                      default=None,
                      help='Key of parent entity to insert into. This should '
                      'be in python tuple-literal form, e.g. ("Foo", 1)')
  return parser


def _GetClient(creds_file, project_id=None, namespace=None):
  """Get a datastore client instance.

  Args:
    creds_file: Path to JSON creds file.
    project_id: Optional project_id of datastore. If not supplied,
                will use one from creds file.
    namespace: Optional namespace to insert into.
  """
  if project_id is None:
    with open(creds_file, 'r') as f:
      project_id = json.load(f)['project_id']

  return datastore.Client.from_service_account_json(
      creds_file, project_id, namespace=namespace), project_id


class DuplicateKeyError(ValueError):
  """Raised when two Entities have the same key."""


def GetEntities(project_id, json_lines, outer_parent_key=None, namespace=None):
  """Create gcloud entities from json string entries.

  project_id: String gcloud project id that entities are for.
  json_lines: File or other line-by-line iterator of json strings to turn into
              entities.
  outer_parent_key: Optional datastore.Key instance to act as the parent_key
                    of all top level entities.
  namespace: Optional string namespace for entities.
  """
  entity_keys = {}

  for line in json_lines:
    item = json.loads(line)
    kind, idx = item.pop('id')
    parent = item.pop('parent', None)

    if (kind, idx) in entity_keys:
      raise DuplicateKeyError(
          'Duplicate entities with id (%s, %s)' % (kind, idx))

    if parent:
      parent_key = entity_keys[tuple(parent)]
    else:
      parent_key = outer_parent_key

    key = datastore.Key(
        kind, idx, project=project_id, parent=parent_key, namespace=namespace)
    e = datastore.Entity(key=key)
    e.update(item)
    entity_keys[(kind, idx)] = key
    entity_keys[idx] = key

    yield e


def main(argv):
  parser = GetParser()
  options = parser.parse_args(argv)

  entities_path = options.entities
  creds_file = options.service_acct_json
  project_id = options.project_id
  namespace = options.namespace

  entities = []
  c, project_id = _GetClient(creds_file, project_id, namespace)

  if options.parent_key:
    upper_parent_key = c.key(*ast.literal_eval(options.parent_key))
  else:
    upper_parent_key = None


  with open(entities_path, 'r') as f:
    entities = GetEntities(project_id, f, upper_parent_key, namespace)

    for chunk in iter_utils.SplitToChunks(entities, _BATCH_CHUNK_SIZE):
      batch = c.batch()
      for e in chunk:
        batch.put(e)
      batch.commit()