# Copyright 2018 Google LLC
#
# Licensed 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.


import unittest

import mock

from ._testing import _make_credentials

from google.cloud.bigtable.batcher import MutationsBatcher
from google.cloud.bigtable.row import DirectRow


class TestMutationsBatcher(unittest.TestCase):
    from grpc import StatusCode

    TABLE_ID = 'table-id'
    TABLE_NAME = '/tables/' + TABLE_ID

    # RPC Status Codes
    SUCCESS = StatusCode.OK.value[0]

    @staticmethod
    def _get_target_class():
        from google.cloud.bigtable.table import Table

        return Table

    def _make_table(self, *args, **kwargs):
        return self._get_target_class()(*args, **kwargs)

    @staticmethod
    def _get_target_client_class():
        from google.cloud.bigtable.client import Client

        return Client

    def _make_client(self, *args, **kwargs):
        return self._get_target_client_class()(*args, **kwargs)

    def test_constructor(self):
        credentials = _make_credentials()
        client = self._make_client(project='project-id',
                                   credentials=credentials, admin=True)

        instance = client.instance(instance_id='instance-id')
        table = self._make_table(self.TABLE_ID, instance)

        mutation_batcher = MutationsBatcher(table)
        self.assertEqual(table, mutation_batcher.table)

    def test_mutate_row(self):
        table = _Table(self.TABLE_NAME)
        mutation_batcher = MutationsBatcher(table=table)

        rows = [DirectRow(row_key=b'row_key'),
                DirectRow(row_key=b'row_key_2'),
                DirectRow(row_key=b'row_key_3'),
                DirectRow(row_key=b'row_key_4')]

        mutation_batcher.mutate_rows(rows)
        mutation_batcher.flush()

        self.assertEqual(table.mutation_calls, 1)

    def test_mutate_rows(self):
        table = _Table(self.TABLE_NAME)
        mutation_batcher = MutationsBatcher(table=table)

        row = DirectRow(row_key=b'row_key')
        row.set_cell('cf1', b'c1', 1)
        row.set_cell('cf1', b'c2', 2)
        row.set_cell('cf1', b'c3', 3)
        row.set_cell('cf1', b'c4', 4)

        mutation_batcher.mutate(row)

        mutation_batcher.flush()

        self.assertEqual(table.mutation_calls, 1)

    def test_flush_with_no_rows(self):
        table = _Table(self.TABLE_NAME)
        mutation_batcher = MutationsBatcher(table=table)
        mutation_batcher.flush()

        self.assertEqual(table.mutation_calls, 0)

    def test_add_row_with_max_flush_count(self):
        table = _Table(self.TABLE_NAME)
        mutation_batcher = MutationsBatcher(table=table, flush_count=3)

        row_1 = DirectRow(row_key=b'row_key_1')
        row_2 = DirectRow(row_key=b'row_key_2')
        row_3 = DirectRow(row_key=b'row_key_3')

        mutation_batcher.mutate(row_1)
        mutation_batcher.mutate(row_2)
        mutation_batcher.mutate(row_3)

        self.assertEqual(table.mutation_calls, 1)

    @mock.patch('google.cloud.bigtable.batcher.MAX_MUTATIONS', new=3)
    def test_mutate_row_with_max_mutations_failure(self):
        from google.cloud.bigtable.batcher import MaxMutationsError

        table = _Table(self.TABLE_NAME)
        mutation_batcher = MutationsBatcher(table=table)

        row = DirectRow(row_key=b'row_key')
        row.set_cell('cf1', b'c1', 1)
        row.set_cell('cf1', b'c2', 2)
        row.set_cell('cf1', b'c3', 3)
        row.set_cell('cf1', b'c4', 4)

        with self.assertRaises(MaxMutationsError):
            mutation_batcher.mutate(row)

    @mock.patch('google.cloud.bigtable.batcher.MAX_MUTATIONS', new=3)
    def test_mutate_row_with_max_mutations(self):
        table = _Table(self.TABLE_NAME)
        mutation_batcher = MutationsBatcher(table=table)

        row = DirectRow(row_key=b'row_key')
        row.set_cell('cf1', b'c1', 1)
        row.set_cell('cf1', b'c2', 2)
        row.set_cell('cf1', b'c3', 3)

        mutation_batcher.mutate(row)
        mutation_batcher.flush()

        self.assertEqual(table.mutation_calls, 1)

    def test_mutate_row_with_max_row_bytes(self):
        table = _Table(self.TABLE_NAME)
        mutation_batcher = MutationsBatcher(table=table,
                                            max_row_bytes=3 * 1024 * 1024)

        number_of_bytes = 1 * 1024 * 1024
        max_value = b'1' * number_of_bytes

        row = DirectRow(row_key=b'row_key')
        row.set_cell('cf1', b'c1', max_value)
        row.set_cell('cf1', b'c2', max_value)
        row.set_cell('cf1', b'c3', max_value)

        mutation_batcher.mutate(row)

        self.assertEqual(table.mutation_calls, 1)


class _Instance(object):

    def __init__(self, client=None):
        self._client = client


class _Table(object):

    def __init__(self, name, client=None):
        self.name = name
        self._instance = _Instance(client)
        self.mutation_calls = 0

    def mutate_rows(self, rows):
        self.mutation_calls += 1
        return rows
