#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Unittest for GCP Bigtable testing."""
from __future__ import absolute_import

import datetime
import logging
import random
import string
import unittest
import uuid

import pytz

import apache_beam as beam
from apache_beam.io.gcp.bigtableio import WriteToBigTable
from apache_beam.metrics.metric import MetricsFilter
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.runners.runner import PipelineState
from apache_beam.testing.test_pipeline import TestPipeline

# Protect against environments where bigtable library is not available.
# pylint: disable=wrong-import-order, wrong-import-position
try:
  from google.cloud._helpers import _datetime_from_microseconds
  from google.cloud._helpers import _microseconds_from_datetime
  from google.cloud._helpers import UTC
  from google.cloud.bigtable import row, column_family, Client
except ImportError:
  Client = None
  UTC = pytz.utc
  _microseconds_from_datetime = lambda label_stamp: label_stamp
  _datetime_from_microseconds = lambda micro: micro


EXISTING_INSTANCES = []
LABEL_KEY = u'python-bigtable-beam'
label_stamp = datetime.datetime.utcnow().replace(tzinfo=UTC)
label_stamp_micros = _microseconds_from_datetime(label_stamp)
LABELS = {LABEL_KEY: str(label_stamp_micros)}


class GenerateTestRows(beam.PTransform):
  """ A transform test to run write to the Bigtable Table.

  A PTransform that generate a list of `DirectRow` to write it in
  Bigtable Table.

  """
  def __init__(self, number, project_id=None, instance_id=None,
               table_id=None):
    # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
    # super(WriteToBigTable, self).__init__()
    beam.PTransform.__init__(self)
    self.number = number
    self.rand = random.choice(string.ascii_letters + string.digits)
    self.column_family_id = 'cf1'
    self.beam_options = {'project_id': project_id,
                         'instance_id': instance_id,
                         'table_id': table_id}

  def _generate(self):
    value = ''.join(self.rand for i in range(100))

    for index in range(self.number):
      key = "beam_key%s" % ('{0:07}'.format(index))
      direct_row = row.DirectRow(row_key=key)
      for column_id in range(10):
        direct_row.set_cell(self.column_family_id,
                            ('field%s' % column_id).encode('utf-8'),
                            value,
                            datetime.datetime.now())
      yield direct_row

  def expand(self, pvalue):
    beam_options = self.beam_options
    return (pvalue
            | beam.Create(self._generate())
            | WriteToBigTable(beam_options['project_id'],
                              beam_options['instance_id'],
                              beam_options['table_id']))


@unittest.skipIf(Client is None, 'GCP Bigtable dependencies are not installed')
class BigtableIOWriteTest(unittest.TestCase):
  """ Bigtable Write Connector Test

  """
  DEFAULT_TABLE_PREFIX = "python-test"
  instance_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
  cluster_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
  table_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
  number = 500
  LOCATION_ID = "us-east1-b"

  def setUp(self):
    try:
      from google.cloud.bigtable import enums
      self.STORAGE_TYPE = enums.StorageType.HDD
      self.INSTANCE_TYPE = enums.Instance.Type.DEVELOPMENT
    except ImportError:
      self.STORAGE_TYPE = 2
      self.INSTANCE_TYPE = 2

    self.test_pipeline = TestPipeline(is_integration_test=True)
    self.runner_name = type(self.test_pipeline.runner).__name__
    self.project = self.test_pipeline.get_option('project')
    self.client = Client(project=self.project, admin=True)

    self._delete_old_instances()

    self.instance = self.client.instance(self.instance_id,
                                         instance_type=self.INSTANCE_TYPE,
                                         labels=LABELS)

    if not self.instance.exists():
      cluster = self.instance.cluster(self.cluster_id,
                                      self.LOCATION_ID,
                                      default_storage_type=self.STORAGE_TYPE)
      self.instance.create(clusters=[cluster])
    self.table = self.instance.table(self.table_id)

    if not self.table.exists():
      max_versions_rule = column_family.MaxVersionsGCRule(2)
      column_family_id = 'cf1'
      column_families = {column_family_id: max_versions_rule}
      self.table.create(column_families=column_families)

  def _delete_old_instances(self):
    instances = self.client.list_instances()
    EXISTING_INSTANCES[:] = instances

    def age_in_hours(micros):
      return (datetime.datetime.utcnow().replace(tzinfo=UTC) - (
          _datetime_from_microseconds(micros))).total_seconds() // 3600
    CLEAN_INSTANCE = [i for instance in EXISTING_INSTANCES for i in instance if(
        LABEL_KEY in i.labels.keys() and
        (age_in_hours(int(i.labels[LABEL_KEY])) >= 2))]

    if CLEAN_INSTANCE:
      for instance in CLEAN_INSTANCE:
        instance.delete()

  def tearDown(self):
    if self.instance.exists():
      self.instance.delete()

  def test_bigtable_write(self):
    number = self.number
    pipeline_args = self.test_pipeline.options_list
    pipeline_options = PipelineOptions(pipeline_args)

    with beam.Pipeline(options=pipeline_options) as pipeline:
      config_data = {'project_id':self.project,
                     'instance_id':self.instance,
                     'table_id':self.table}
      _ = (
          pipeline
          | 'Generate Direct Rows' >> GenerateTestRows(number, **config_data))

      result = pipeline.run()
      result.wait_until_finish()

      assert result.state == PipelineState.DONE

      read_rows = self.table.read_rows()
      assert len([_ for _ in read_rows]) == number

      if not hasattr(result, 'has_job') or result.has_job:
        read_filter = MetricsFilter().with_name('Written Row')
        query_result = result.metrics().query(read_filter)
        if query_result['counters']:
          read_counter = query_result['counters'][0]

          logging.info('Number of Rows: %d', read_counter.committed)
          assert read_counter.committed == number


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()