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

"""Tests common to all coder implementations."""
# pytype: skip-file

from __future__ import absolute_import

import logging
import math
import sys
import unittest
from builtins import range

import pytest

from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message
from apache_beam.coders import coders
from apache_beam.internal import pickler
from apache_beam.runners import pipeline_context
from apache_beam.transforms import window
from apache_beam.transforms.window import GlobalWindow
from apache_beam.utils import timestamp
from apache_beam.utils import windowed_value
from apache_beam.utils.timestamp import MIN_TIMESTAMP

from . import observable


# Defined out of line for picklability.
class CustomCoder(coders.Coder):

  def encode(self, x):
    return str(x+1).encode('utf-8')

  def decode(self, encoded):
    return int(encoded) - 1


# These tests need to all be run in the same process due to the asserts
# in tearDownClass.
@pytest.mark.no_xdist
class CodersTest(unittest.TestCase):

  # These class methods ensure that we test each defined coder in both
  # nested and unnested context.

  @classmethod
  def setUpClass(cls):
    cls.seen = set()
    cls.seen_nested = set()
    # Method has been renamed in Python 3
    if sys.version_info[0] < 3:
      cls.assertCountEqual = cls.assertItemsEqual

  @classmethod
  def tearDownClass(cls):
    standard = set(c
                   for c in coders.__dict__.values()
                   if isinstance(c, type) and issubclass(c, coders.Coder) and
                   'Base' not in c.__name__)
    standard -= set([coders.Coder,
                     coders.AvroGenericCoder,
                     coders.DeterministicProtoCoder,
                     coders.FastCoder,
                     coders.ProtoCoder,
                     coders.RunnerAPICoderHolder,
                     coders.ToStringCoder])
    assert not standard - cls.seen, standard - cls.seen
    assert not standard - cls.seen_nested, standard - cls.seen_nested

  @classmethod
  def _observe(cls, coder):
    cls.seen.add(type(coder))
    cls._observe_nested(coder)

  @classmethod
  def _observe_nested(cls, coder):
    if isinstance(coder, coders.TupleCoder):
      for c in coder.coders():
        cls.seen_nested.add(type(c))
        cls._observe_nested(c)

  def check_coder(self, coder, *values, **kwargs):
    context = kwargs.pop('context', pipeline_context.PipelineContext())
    test_size_estimation = kwargs.pop('test_size_estimation', True)
    assert not kwargs
    self._observe(coder)
    for v in values:
      self.assertEqual(v, coder.decode(coder.encode(v)))
      if test_size_estimation:
        self.assertEqual(coder.estimate_size(v),
                         len(coder.encode(v)))
        self.assertEqual(coder.estimate_size(v),
                         coder.get_impl().estimate_size(v))
        self.assertEqual(coder.get_impl().get_estimated_size_and_observables(v),
                         (coder.get_impl().estimate_size(v), []))
      copy1 = pickler.loads(pickler.dumps(coder))
    copy2 = coders.Coder.from_runner_api(coder.to_runner_api(context), context)
    for v in values:
      self.assertEqual(v, copy1.decode(copy2.encode(v)))
      if coder.is_deterministic():
        self.assertEqual(copy1.encode(v), copy2.encode(v))

  def test_custom_coder(self):

    self.check_coder(CustomCoder(), 1, -10, 5)
    self.check_coder(coders.TupleCoder((CustomCoder(), coders.BytesCoder())),
                     (1, b'a'), (-10, b'b'), (5, b'c'))

  def test_pickle_coder(self):
    self.check_coder(coders.PickleCoder(), 'a', 1, 1.5, (1, 2, 3))

  def test_deterministic_coder(self):
    coder = coders.FastPrimitivesCoder()
    deterministic_coder = coders.DeterministicFastPrimitivesCoder(coder, 'step')
    self.check_coder(deterministic_coder, 'a', 1, 1.5, (1, 2, 3))
    with self.assertRaises(TypeError):
      self.check_coder(deterministic_coder, dict())
    with self.assertRaises(TypeError):
      self.check_coder(deterministic_coder, [1, dict()])

    self.check_coder(coders.TupleCoder((deterministic_coder, coder)),
                     (1, dict()), ('a', [dict()]))

  def test_dill_coder(self):
    cell_value = (lambda x: lambda: x)(0).__closure__[0]
    self.check_coder(coders.DillCoder(), 'a', 1, cell_value)
    self.check_coder(
        coders.TupleCoder((coders.VarIntCoder(), coders.DillCoder())),
        (1, cell_value))

  def test_fast_primitives_coder(self):
    coder = coders.FastPrimitivesCoder(coders.SingletonCoder(len))
    self.check_coder(coder, None, 1, -1, 1.5, b'str\0str', u'unicode\0\u0101')
    self.check_coder(coder, (), (1, 2, 3))
    self.check_coder(coder, [], [1, 2, 3])
    self.check_coder(coder, dict(), {'a': 'b'}, {0: dict(), 1: len})
    self.check_coder(coder, set(), {'a', 'b'})
    self.check_coder(coder, True, False)
    self.check_coder(coder, len)
    self.check_coder(coders.TupleCoder((coder,)), ('a',), (1,))

  def test_fast_primitives_coder_large_int(self):
    coder = coders.FastPrimitivesCoder()
    self.check_coder(coder, 10 ** 100)

  def test_bytes_coder(self):
    self.check_coder(coders.BytesCoder(), b'a', b'\0', b'z' * 1000)

  def test_bool_coder(self):
    self.check_coder(coders.BooleanCoder(), True, False)

  def test_varint_coder(self):
    # Small ints.
    self.check_coder(coders.VarIntCoder(), *range(-10, 10))
    # Multi-byte encoding starts at 128
    self.check_coder(coders.VarIntCoder(), *range(120, 140))
    # Large values
    MAX_64_BIT_INT = 0x7fffffffffffffff
    self.check_coder(coders.VarIntCoder(),
                     *[int(math.pow(-1, k) * math.exp(k))
                       for k in range(0, int(math.log(MAX_64_BIT_INT)))])

  def test_float_coder(self):
    self.check_coder(coders.FloatCoder(),
                     *[float(0.1 * x) for x in range(-100, 100)])
    self.check_coder(coders.FloatCoder(),
                     *[float(2 ** (0.1 * x)) for x in range(-100, 100)])
    self.check_coder(coders.FloatCoder(), float('-Inf'), float('Inf'))
    self.check_coder(
        coders.TupleCoder((coders.FloatCoder(), coders.FloatCoder())),
        (0, 1), (-100, 100), (0.5, 0.25))

  def test_singleton_coder(self):
    a = 'anything'
    b = 'something else'
    self.check_coder(coders.SingletonCoder(a), a)
    self.check_coder(coders.SingletonCoder(b), b)
    self.check_coder(coders.TupleCoder((coders.SingletonCoder(a),
                                        coders.SingletonCoder(b))), (a, b))

  def test_interval_window_coder(self):
    self.check_coder(coders.IntervalWindowCoder(),
                     *[window.IntervalWindow(x, y)
                       for x in [-2**52, 0, 2**52]
                       for y in range(-100, 100)])
    self.check_coder(
        coders.TupleCoder((coders.IntervalWindowCoder(),)),
        (window.IntervalWindow(0, 10),))

  def test_timestamp_coder(self):
    self.check_coder(coders.TimestampCoder(),
                     *[timestamp.Timestamp(micros=x) for x in (-1000, 0, 1000)])
    self.check_coder(coders.TimestampCoder(),
                     timestamp.Timestamp(micros=-1234567000),
                     timestamp.Timestamp(micros=1234567000))
    self.check_coder(coders.TimestampCoder(),
                     timestamp.Timestamp(micros=-1234567890123456000),
                     timestamp.Timestamp(micros=1234567890123456000))
    self.check_coder(
        coders.TupleCoder((coders.TimestampCoder(), coders.BytesCoder())),
        (timestamp.Timestamp.of(27), b'abc'))

  def test_timer_coder(self):
    self.check_coder(coders._TimerCoder(coders.BytesCoder()),
                     *[{'timestamp': timestamp.Timestamp(micros=x),
                        'payload': b'xyz'}
                       for x in (-3000, 0, 3000)])
    self.check_coder(
        coders.TupleCoder((coders._TimerCoder(coders.VarIntCoder()),)),
        ({'timestamp': timestamp.Timestamp.of(37000), 'payload': 389},))

  def test_tuple_coder(self):
    kv_coder = coders.TupleCoder((coders.VarIntCoder(), coders.BytesCoder()))
    # Verify cloud object representation
    self.assertEqual(
        {
            '@type': 'kind:pair',
            'is_pair_like': True,
            'component_encodings': [
                coders.VarIntCoder().as_cloud_object(),
                coders.BytesCoder().as_cloud_object()],
        },
        kv_coder.as_cloud_object())
    # Test binary representation
    self.assertEqual(
        b'\x04abc',
        kv_coder.encode((4, b'abc')))
    # Test unnested
    self.check_coder(
        kv_coder,
        (1, b'a'),
        (-2, b'a' * 100),
        (300, b'abc\0' * 5))
    # Test nested
    self.check_coder(
        coders.TupleCoder(
            (coders.TupleCoder((coders.PickleCoder(), coders.VarIntCoder())),
             coders.StrUtf8Coder(),
             coders.BooleanCoder())),
        ((1, 2), 'a', True),
        ((-2, 5), u'a\u0101' * 100, False),
        ((300, 1), 'abc\0' * 5, True))

  def test_tuple_sequence_coder(self):
    int_tuple_coder = coders.TupleSequenceCoder(coders.VarIntCoder())
    self.check_coder(int_tuple_coder, (1, -1, 0), (), tuple(range(1000)))
    self.check_coder(
        coders.TupleCoder((coders.VarIntCoder(), int_tuple_coder)),
        (1, (1, 2, 3)))

  def test_base64_pickle_coder(self):
    self.check_coder(coders.Base64PickleCoder(), 'a', 1, 1.5, (1, 2, 3))

  def test_utf8_coder(self):
    self.check_coder(coders.StrUtf8Coder(), 'a', u'ab\u00FF', u'\u0101\0')

  def test_iterable_coder(self):
    iterable_coder = coders.IterableCoder(coders.VarIntCoder())
    # Verify cloud object representation
    self.assertEqual(
        {
            '@type': 'kind:stream',
            'is_stream_like': True,
            'component_encodings': [coders.VarIntCoder().as_cloud_object()]
        },
        iterable_coder.as_cloud_object())
    # Test unnested
    self.check_coder(iterable_coder,
                     [1], [-1, 0, 100])
    # Test nested
    self.check_coder(
        coders.TupleCoder((coders.VarIntCoder(),
                           coders.IterableCoder(coders.VarIntCoder()))),
        (1, [1, 2, 3]))

  def test_iterable_coder_unknown_length(self):
    # Empty
    self._test_iterable_coder_of_unknown_length(0)
    # Single element
    self._test_iterable_coder_of_unknown_length(1)
    # Multiple elements
    self._test_iterable_coder_of_unknown_length(100)
    # Multiple elements with underlying stream buffer overflow.
    self._test_iterable_coder_of_unknown_length(80000)

  def _test_iterable_coder_of_unknown_length(self, count):
    def iter_generator(count):
      for i in range(count):
        yield i

    iterable_coder = coders.IterableCoder(coders.VarIntCoder())
    self.assertCountEqual(list(iter_generator(count)),
                          iterable_coder.decode(
                              iterable_coder.encode(iter_generator(count))))

  def test_windowedvalue_coder_paneinfo(self):
    coder = coders.WindowedValueCoder(coders.VarIntCoder(),
                                      coders.GlobalWindowCoder())
    test_paneinfo_values = [
        windowed_value.PANE_INFO_UNKNOWN,
        windowed_value.PaneInfo(
            True, True, windowed_value.PaneInfoTiming.EARLY, 0, -1),
        windowed_value.PaneInfo(
            True, False, windowed_value.PaneInfoTiming.ON_TIME, 0, 0),
        windowed_value.PaneInfo(
            True, False, windowed_value.PaneInfoTiming.ON_TIME, 10, 0),
        windowed_value.PaneInfo(
            False, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 23),
        windowed_value.PaneInfo(
            False, True, windowed_value.PaneInfoTiming.ON_TIME, 12, 23),
        windowed_value.PaneInfo(
            False, False, windowed_value.PaneInfoTiming.LATE, 0, 123),]

    test_values = [windowed_value.WindowedValue(123, 234, (GlobalWindow(),), p)
                   for p in test_paneinfo_values]

    # Test unnested.
    self.check_coder(coder, windowed_value.WindowedValue(
        123, 234, (GlobalWindow(),), windowed_value.PANE_INFO_UNKNOWN))
    for value in test_values:
      self.check_coder(coder, value)

    # Test nested.
    for value1 in test_values:
      for value2 in test_values:
        self.check_coder(coders.TupleCoder((coder, coder)), (value1, value2))

  def test_windowed_value_coder(self):
    coder = coders.WindowedValueCoder(coders.VarIntCoder(),
                                      coders.GlobalWindowCoder())
    # Verify cloud object representation
    self.assertEqual(
        {
            '@type': 'kind:windowed_value',
            'is_wrapper': True,
            'component_encodings': [
                coders.VarIntCoder().as_cloud_object(),
                coders.GlobalWindowCoder().as_cloud_object(),
            ],
        },
        coder.as_cloud_object())
    # Test binary representation
    self.assertEqual(b'\x7f\xdf;dZ\x1c\xac\t\x00\x00\x00\x01\x0f\x01',
                     coder.encode(window.GlobalWindows.windowed_value(1)))

    # Test decoding large timestamp
    self.assertEqual(
        coder.decode(b'\x7f\xdf;dZ\x1c\xac\x08\x00\x00\x00\x01\x0f\x00'),
        windowed_value.create(0, MIN_TIMESTAMP.micros, (GlobalWindow(),)))

    # Test unnested
    self.check_coder(
        coders.WindowedValueCoder(coders.VarIntCoder()),
        windowed_value.WindowedValue(3, -100, ()),
        windowed_value.WindowedValue(-1, 100, (1, 2, 3)))

    # Test Global Window
    self.check_coder(
        coders.WindowedValueCoder(coders.VarIntCoder(),
                                  coders.GlobalWindowCoder()),
        window.GlobalWindows.windowed_value(1))

    # Test nested
    self.check_coder(
        coders.TupleCoder((
            coders.WindowedValueCoder(coders.FloatCoder()),
            coders.WindowedValueCoder(coders.StrUtf8Coder()))),
        (windowed_value.WindowedValue(1.5, 0, ()),
         windowed_value.WindowedValue("abc", 10, ('window',))))

  def test_param_windowed_value_coder(self):
    from apache_beam.transforms.window import IntervalWindow
    from apache_beam.utils.windowed_value import PaneInfo
    wv = windowed_value.create(
        b'',
        # Milliseconds to microseconds
        1000 * 1000,
        (IntervalWindow(11, 21),),
        PaneInfo(True, False, 1, 2, 3))
    windowed_value_coder = coders.WindowedValueCoder(
        coders.BytesCoder(), coders.IntervalWindowCoder())
    payload = windowed_value_coder.encode(wv)
    coder = coders.ParamWindowedValueCoder(
        payload, [coders.VarIntCoder(), coders.IntervalWindowCoder()])

    # Test binary representation
    self.assertEqual(b'\x01',
                     coder.encode(window.GlobalWindows.windowed_value(1)))

    # Test unnested
    self.check_coder(
        coders.ParamWindowedValueCoder(
            payload, [coders.VarIntCoder(), coders.IntervalWindowCoder()]),
        windowed_value.WindowedValue(
            3,
            1,
            (window.IntervalWindow(11, 21),),
            PaneInfo(True, False, 1, 2, 3)),
        windowed_value.WindowedValue(
            1,
            1,
            (window.IntervalWindow(11, 21),),
            PaneInfo(True, False, 1, 2, 3)))

    # Test nested
    self.check_coder(
        coders.TupleCoder((
            coders.ParamWindowedValueCoder(
                payload, [
                    coders.FloatCoder(),
                    coders.IntervalWindowCoder()]),
            coders.ParamWindowedValueCoder(
                payload, [
                    coders.StrUtf8Coder(),
                    coders.IntervalWindowCoder()]))),
        (windowed_value.WindowedValue(
            1.5,
            1,
            (window.IntervalWindow(11, 21),),
            PaneInfo(True, False, 1, 2, 3)),
         windowed_value.WindowedValue(
             "abc",
             1,
             (window.IntervalWindow(11, 21),),
             PaneInfo(True, False, 1, 2, 3))))

  def test_proto_coder(self):
    # For instructions on how these test proto message were generated,
    # see coders_test.py
    ma = test_message.MessageA()
    mab = ma.field2.add()
    mab.field1 = True
    ma.field1 = u'hello world'

    mb = test_message.MessageA()
    mb.field1 = u'beam'

    proto_coder = coders.ProtoCoder(ma.__class__)
    self.check_coder(proto_coder, ma)
    self.check_coder(coders.TupleCoder((proto_coder, coders.BytesCoder())),
                     (ma, b'a'), (mb, b'b'))

  def test_global_window_coder(self):
    coder = coders.GlobalWindowCoder()
    value = window.GlobalWindow()
    # Verify cloud object representation
    self.assertEqual({'@type': 'kind:global_window'},
                     coder.as_cloud_object())
    # Test binary representation
    self.assertEqual(b'', coder.encode(value))
    self.assertEqual(value, coder.decode(b''))
    # Test unnested
    self.check_coder(coder, value)
    # Test nested
    self.check_coder(coders.TupleCoder((coder, coder)),
                     (value, value))

  def test_length_prefix_coder(self):
    coder = coders.LengthPrefixCoder(coders.BytesCoder())
    # Verify cloud object representation
    self.assertEqual(
        {
            '@type': 'kind:length_prefix',
            'component_encodings': [coders.BytesCoder().as_cloud_object()]
        },
        coder.as_cloud_object())
    # Test binary representation
    self.assertEqual(b'\x00', coder.encode(b''))
    self.assertEqual(b'\x01a', coder.encode(b'a'))
    self.assertEqual(b'\x02bc', coder.encode(b'bc'))
    self.assertEqual(b'\xff\x7f' + b'z' * 16383, coder.encode(b'z' * 16383))
    # Test unnested
    self.check_coder(coder, b'', b'a', b'bc', b'def')
    # Test nested
    self.check_coder(coders.TupleCoder((coder, coder)),
                     (b'', b'a'),
                     (b'bc', b'def'))

  def test_nested_observables(self):
    class FakeObservableIterator(observable.ObservableMixin):

      def __iter__(self):
        return iter([1, 2, 3])

    # Coder for elements from the observable iterator.
    elem_coder = coders.VarIntCoder()
    iter_coder = coders.TupleSequenceCoder(elem_coder)

    # Test nested WindowedValue observable.
    coder = coders.WindowedValueCoder(iter_coder)
    observ = FakeObservableIterator()
    value = windowed_value.WindowedValue(observ, 0, ())
    self.assertEqual(
        coder.get_impl().get_estimated_size_and_observables(value)[1],
        [(observ, elem_coder.get_impl())])

    # Test nested tuple observable.
    coder = coders.TupleCoder((coders.StrUtf8Coder(), iter_coder))
    value = (u'123', observ)
    self.assertEqual(
        coder.get_impl().get_estimated_size_and_observables(value)[1],
        [(observ, elem_coder.get_impl())])

  def test_state_backed_iterable_coder(self):
    # pylint: disable=global-variable-undefined
    # required for pickling by reference
    global state
    state = {}

    def iterable_state_write(values, element_coder_impl):
      token = b'state_token_%d' % len(state)
      state[token] = [element_coder_impl.encode(e) for e in values]
      return token

    def iterable_state_read(token, element_coder_impl):
      return [element_coder_impl.decode(s) for s in state[token]]

    coder = coders.StateBackedIterableCoder(
        coders.VarIntCoder(),
        read_state=iterable_state_read,
        write_state=iterable_state_write,
        write_state_threshold=1)
    context = pipeline_context.PipelineContext(
        iterable_state_read=iterable_state_read,
        iterable_state_write=iterable_state_write)
    self.check_coder(
        coder, [1, 2, 3], context=context, test_size_estimation=False)
    # Ensure that state was actually used.
    self.assertNotEqual(state, {})
    self.check_coder(coders.TupleCoder((coder, coder)),
                     ([1], [2, 3]),
                     context=context,
                     test_size_estimation=False)


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