#
# 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 userstate
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.ExternalCoder,
        coders.FastCoder,
        coders.ProtoCoder,
        coders.ToBytesCoder
    ])
    cls.seen_nested -= set([coders.ProtoCoder, CustomCoder])
    assert not standard - cls.seen
    assert not cls.seen_nested - standard

  @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.StrUtf8Coder(), coders.GlobalWindowCoder()),
        *[
            userstate.Timer(
                user_key="key",
                dynamic_timer_tag="tag",
                windows=(GlobalWindow(), ),
                clear_bit=True,
                fire_timestamp=None,
                hold_timestamp=None,
                paneinfo=None),
            userstate.Timer(
                user_key="key",
                dynamic_timer_tag="tag",
                windows=(GlobalWindow(), ),
                clear_bit=False,
                fire_timestamp=timestamp.Timestamp.of(123),
                hold_timestamp=timestamp.Timestamp.of(456),
                paneinfo=windowed_value.PANE_INFO_UNKNOWN)
        ])

  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()