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

import copy
import dataclasses
import logging
import os
import unittest
from typing import Optional

import pytest
from parameterized import parameterized

from apache_beam.internal.cloudpickle import cloudpickle
from apache_beam.ml.anomaly.specifiable import _FALLBACK_SUBSPACE
from apache_beam.ml.anomaly.specifiable import _KNOWN_SPECIFIABLE
from apache_beam.ml.anomaly.specifiable import Spec
from apache_beam.ml.anomaly.specifiable import Specifiable
from apache_beam.ml.anomaly.specifiable import specifiable


class TestSpecifiable(unittest.TestCase):
  def setUp(self) -> None:
    self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE)

  def tearDown(self) -> None:
    _KNOWN_SPECIFIABLE.clear()
    _KNOWN_SPECIFIABLE.update(self.saved_specifiable)

  def test_decorator_in_function_form(self):
    class A():
      pass

    class B():
      pass

    # class is not decorated and thus not registered
    self.assertNotIn("A", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE])

    # apply the decorator function to an existing class
    A = specifiable(A)
    self.assertEqual(A.spec_type(), "A")
    self.assertTrue(isinstance(A(), Specifiable))
    self.assertIn("A", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE])
    self.assertEqual(_KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]["A"], A)

    # Re-registering spec_type with the same class is allowed
    A = specifiable(A)

    # Raise an error when re-registering spec_type with a different class
    self.assertRaises(ValueError, specifiable(spec_type='A'), B)

    # Applying the decorator function to an existing class with a different
    # spec_type will have no effect.
    A = specifiable(spec_type="A_DUP")(A)
    self.assertEqual(A.spec_type(), "A")

  def test_decorator_in_syntactic_sugar_form(self):
    # call decorator without parameters
    @specifiable
    class B():
      pass

    self.assertTrue(isinstance(B(), Specifiable))
    self.assertIn("B", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE])
    self.assertEqual(_KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]["B"], B)

    # call decorator with parameters
    @specifiable(spec_type="C_TYPE")
    class C():
      pass

    self.assertTrue(isinstance(C(), Specifiable))
    self.assertIn("C_TYPE", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE])
    self.assertEqual(_KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]["C_TYPE"], C)

  def test_init_params_in_specifiable(self):
    @specifiable
    class ParentWithInitParams():
      def __init__(self, arg_1, arg_2=2, arg_3="3", **kwargs):
        pass

    parent = ParentWithInitParams(10, arg_3="30", arg_4=40)
    assert isinstance(parent, Specifiable)
    self.assertEqual(
        parent.init_kwargs, {
            'arg_1': 10, 'arg_3': '30', 'arg_4': 40
        })

    # inheritance of a Specifiable subclass
    @specifiable
    class ChildWithInitParams(ParentWithInitParams):
      def __init__(self, new_arg_1, new_arg_2=200, new_arg_3="300", **kwargs):
        super().__init__(**kwargs)

    child = ChildWithInitParams(
        1000, arg_1=11, arg_2=20, new_arg_2=2000, arg_4=4000)
    assert isinstance(child, Specifiable)
    self.assertEqual(
        child.init_kwargs,
        {
            'new_arg_1': 1000,
            'arg_1': 11,
            'arg_2': 20,
            'new_arg_2': 2000,
            'arg_4': 4000
        })

    # composite of Specifiable subclasses
    @specifiable
    class CompositeWithInitParams():
      def __init__(
          self,
          my_parent: Optional[ParentWithInitParams] = None,
          my_child: Optional[ChildWithInitParams] = None):
        pass

    composite = CompositeWithInitParams(parent, child)
    assert isinstance(composite, Specifiable)
    self.assertEqual(
        composite.init_kwargs, {
            'my_parent': parent, 'my_child': child
        })

  def test_from_spec_on_unknown_spec_type(self):
    self.assertRaises(ValueError, Specifiable.from_spec, Spec(type="unknown"))

  # To test from_spec and to_spec with/without just_in_time_init.
  @parameterized.expand([(False, False), (True, False), (False, True),
                         (True, True)])
  def test_from_spec_and_to_spec(self, on_demand_init, just_in_time_init):
    @specifiable(
        on_demand_init=on_demand_init, just_in_time_init=just_in_time_init)
    @dataclasses.dataclass
    class Product():
      name: str
      price: float

    @specifiable(
        on_demand_init=on_demand_init, just_in_time_init=just_in_time_init)
    class Entry():
      def __init__(self, product: Product, quantity: int = 1):
        self._product = product
        self._quantity = quantity

      def __eq__(self, value) -> bool:
        return isinstance(value, Entry) and \
          self._product == value._product and \
          self._quantity == value._quantity

    @specifiable(
        on_demand_init=on_demand_init, just_in_time_init=just_in_time_init)
    @dataclasses.dataclass
    class ShoppingCart():
      user_id: str
      entries: list[Entry]

    orange = Product("orange", 1.0)

    expected_orange_spec = Spec(
        "Product", config={
            'name': 'orange', 'price': 1.0
        })
    assert isinstance(orange, Specifiable)
    self.assertEqual(orange.to_spec(), expected_orange_spec)

    entry_1 = Entry(product=orange)

    expected_entry_spec_1 = Spec(
        "Entry", config={
            'product': expected_orange_spec,
        })

    assert isinstance(entry_1, Specifiable)
    self.assertEqual(entry_1.to_spec(), expected_entry_spec_1)

    banana = Product("banana", 0.5)
    expected_banana_spec = Spec(
        "Product", config={
            'name': 'banana', 'price': 0.5
        })
    entry_2 = Entry(product=banana, quantity=5)
    expected_entry_spec_2 = Spec(
        "Entry", config={
            'product': expected_banana_spec, 'quantity': 5
        })

    shopping_cart = ShoppingCart(user_id="test", entries=[entry_1, entry_2])
    expected_shopping_cart_spec = Spec(
        "ShoppingCart",
        config={
            "user_id": "test",
            "entries": [expected_entry_spec_1, expected_entry_spec_2]
        })

    assert isinstance(shopping_cart, Specifiable)
    self.assertEqual(shopping_cart.to_spec(), expected_shopping_cart_spec)
    if on_demand_init and not just_in_time_init:
      orange.run_original_init()
      banana.run_original_init()
      entry_1.run_original_init()
      entry_2.run_original_init()
      shopping_cart.run_original_init()

    self.assertEqual(Specifiable.from_spec(expected_orange_spec), orange)
    self.assertEqual(Specifiable.from_spec(expected_entry_spec_1), entry_1)
    self.assertEqual(
        Specifiable.from_spec(expected_shopping_cart_spec), shopping_cart)


class TestInitCallCount(unittest.TestCase):
  def setUp(self) -> None:
    self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE)

  def tearDown(self) -> None:
    _KNOWN_SPECIFIABLE.clear()
    _KNOWN_SPECIFIABLE.update(self.saved_specifiable)

  def test_on_demand_init(self):
    @specifiable(on_demand_init=True, just_in_time_init=False)
    class FooOnDemand():
      counter = 0

      def __init__(self, arg):
        self.my_arg = arg * 10
        FooOnDemand.counter += 1  # increment it when __init__ is called

    foo = FooOnDemand(123)
    self.assertEqual(FooOnDemand.counter, 0)
    self.assertIn("init_kwargs", foo.__dict__)
    self.assertEqual(foo.__dict__["init_kwargs"], {"arg": 123})

    self.assertNotIn("my_arg", foo.__dict__)
    self.assertRaises(AttributeError, getattr, foo, "my_arg")
    self.assertRaises(AttributeError, lambda: foo.my_arg)
    self.assertRaises(AttributeError, getattr, foo, "unknown_arg")
    self.assertRaises(AttributeError, lambda: foo.unknown_arg)
    self.assertEqual(FooOnDemand.counter, 0)

    # __init__ is called when _run_init=True is used
    foo_2 = FooOnDemand(456, _run_init=True)
    self.assertEqual(FooOnDemand.counter, 1)
    self.assertIn("init_kwargs", foo_2.__dict__)
    self.assertEqual(foo_2.__dict__["init_kwargs"], {"arg": 456})

    self.assertIn("my_arg", foo_2.__dict__)
    self.assertEqual(foo_2.my_arg, 4560)
    self.assertEqual(FooOnDemand.counter, 1)

  def test_just_in_time_init(self):
    @specifiable(on_demand_init=False, just_in_time_init=True)
    class FooJustInTime():
      counter = 0

      def __init__(self, arg):
        self.my_arg = arg * 10
        FooJustInTime.counter += 1  # increment it when __init__ is called

    foo = FooJustInTime(321)
    self.assertEqual(FooJustInTime.counter, 0)
    self.assertIn("init_kwargs", foo.__dict__)
    self.assertEqual(foo.__dict__["init_kwargs"], {"arg": 321})

    # __init__ hasn't been called yet
    self.assertNotIn("my_arg", foo.__dict__)
    self.assertEqual(FooJustInTime.counter, 0)

    # __init__ is called when trying to access a class attribute
    self.assertEqual(foo.my_arg, 3210)
    self.assertEqual(FooJustInTime.counter, 1)
    self.assertRaises(AttributeError, lambda: foo.unknown_arg)
    self.assertEqual(FooJustInTime.counter, 1)

  def test_on_demand_and_just_in_time_init(self):
    @specifiable(on_demand_init=True, just_in_time_init=True)
    class FooOnDemandAndJustInTime():
      counter = 0

      def __init__(self, arg):
        self.my_arg = arg * 10
        FooOnDemandAndJustInTime.counter += 1

    foo = FooOnDemandAndJustInTime(987)
    self.assertEqual(FooOnDemandAndJustInTime.counter, 0)
    self.assertIn("init_kwargs", foo.__dict__)
    self.assertEqual(foo.__dict__["init_kwargs"], {"arg": 987})
    self.assertNotIn("my_arg", foo.__dict__)

    self.assertEqual(FooOnDemandAndJustInTime.counter, 0)
    # __init__ is called when trying to access a class attribute
    self.assertEqual(foo.my_arg, 9870)
    self.assertEqual(FooOnDemandAndJustInTime.counter, 1)

    # __init__ is called when _run_init=True is used
    foo_2 = FooOnDemandAndJustInTime(789, _run_init=True)
    self.assertEqual(FooOnDemandAndJustInTime.counter, 2)
    self.assertIn("init_kwargs", foo_2.__dict__)
    self.assertEqual(foo_2.__dict__["init_kwargs"], {"arg": 789})

    self.assertEqual(FooOnDemandAndJustInTime.counter, 2)
    # __init__ is NOT called after it is initialized
    self.assertEqual(foo_2.my_arg, 7890)
    self.assertEqual(FooOnDemandAndJustInTime.counter, 2)

  @specifiable(on_demand_init=True, just_in_time_init=True)
  class FooForPickle():
    counter = 0

    def __init__(self, arg):
      self.my_arg = arg * 10
      type(self).counter += 1

  @pytest.mark.uses_dill
  def test_on_dill_pickle(self):
    pytest.importorskip("dill")

    FooForPickle = TestInitCallCount.FooForPickle

    import dill
    FooForPickle.counter = 0
    foo = FooForPickle(456)
    self.assertEqual(FooForPickle.counter, 0)
    new_foo = dill.loads(dill.dumps(foo))
    self.assertEqual(FooForPickle.counter, 0)
    self.assertEqual(new_foo.__dict__, foo.__dict__)
    self.assertEqual(foo.my_arg, 4560)
    self.assertEqual(FooForPickle.counter, 1)
    new_foo_2 = dill.loads(dill.dumps(foo))
    self.assertEqual(FooForPickle.counter, 1)
    self.assertEqual(new_foo_2.__dict__, foo.__dict__)

  def test_on_pickle(self):
    FooForPickle = TestInitCallCount.FooForPickle

    # Note that pickle does not support classes/functions nested in a function.
    import pickle
    FooForPickle.counter = 0
    foo = FooForPickle(456)
    self.assertEqual(FooForPickle.counter, 0)
    new_foo = pickle.loads(pickle.dumps(foo))
    self.assertEqual(FooForPickle.counter, 0)
    self.assertEqual(new_foo.__dict__, foo.__dict__)
    self.assertEqual(foo.my_arg, 4560)
    self.assertEqual(FooForPickle.counter, 1)
    new_foo_2 = pickle.loads(pickle.dumps(foo))
    self.assertEqual(FooForPickle.counter, 1)
    self.assertEqual(new_foo_2.__dict__, foo.__dict__)

    FooForPickle.counter = 0
    foo = FooForPickle(456)
    self.assertEqual(FooForPickle.counter, 0)
    new_foo = cloudpickle.loads(cloudpickle.dumps(foo))
    self.assertEqual(FooForPickle.counter, 0)
    self.assertEqual(new_foo.__dict__, foo.__dict__)
    self.assertEqual(foo.my_arg, 4560)
    self.assertEqual(FooForPickle.counter, 1)
    new_foo_2 = cloudpickle.loads(cloudpickle.dumps(foo))
    self.assertEqual(FooForPickle.counter, 1)
    self.assertEqual(new_foo_2.__dict__, foo.__dict__)


@specifiable
class Parent():
  counter = 0
  parent_class_var = 1000

  def __init__(self, p):
    self.parent_inst_var = p * 10
    Parent.counter += 1


@specifiable
class Child_1(Parent):
  counter = 0
  child_class_var = 2001

  def __init__(self, c):
    super().__init__(c)
    self.child_inst_var = c + 1
    Child_1.counter += 1


@specifiable
class Child_2(Parent):
  counter = 0
  child_class_var = 2001

  def __init__(self, c):
    self.child_inst_var = c + 1
    super().__init__(c)
    Child_2.counter += 1


@specifiable
class Child_Error_1(Parent):
  counter = 0
  child_class_var = 2001

  def __init__(self, c):
    # read an instance var in child that doesn't exist
    self.child_inst_var += 1
    super().__init__(c)
    Child_2.counter += 1


@specifiable
class Child_Error_2(Parent):
  counter = 0
  child_class_var = 2001

  def __init__(self, c):
    # read an instance var in parent without calling parent's __init__.
    self.parent_inst_var += 1
    Child_2.counter += 1


class TestNestedSpecifiable(unittest.TestCase):
  def setUp(self) -> None:
    self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE)

  def tearDown(self) -> None:
    _KNOWN_SPECIFIABLE.clear()
    _KNOWN_SPECIFIABLE.update(self.saved_specifiable)

  @parameterized.expand([[Child_1, 0], [Child_2, 0], [Child_1, 1], [Child_2, 1],
                         [Child_1, 2], [Child_2, 2]])
  def test_nested_specifiable(self, Child, mode):
    Parent.counter = 0
    Child.counter = 0
    child = Child(5)

    self.assertEqual(Parent.counter, 0)
    self.assertEqual(Child.counter, 0)

    # accessing class vars won't trigger __init__
    self.assertEqual(child.parent_class_var, 1000)
    self.assertEqual(child.child_class_var, 2001)
    self.assertEqual(Parent.counter, 0)
    self.assertEqual(Child.counter, 0)

    # accessing instance var will trigger __init__
    if mode == 0:
      self.assertEqual(child.parent_inst_var, 50)
    elif mode == 1:
      self.assertEqual(child.child_inst_var, 6)
    else:
      self.assertRaises(AttributeError, lambda: child.unknown_var)

    self.assertEqual(Parent.counter, 1)
    self.assertEqual(Child.counter, 1)

    # after initialization, it won't trigger __init__ again
    self.assertEqual(child.parent_inst_var, 50)
    self.assertEqual(child.child_inst_var, 6)
    self.assertRaises(AttributeError, lambda: child.unknown_var)

    self.assertEqual(Parent.counter, 1)
    self.assertEqual(Child.counter, 1)

  def test_error_in_child(self):
    Parent.counter = 0
    child_1 = Child_Error_1(5)

    self.assertEqual(child_1.child_class_var, 2001)

    # error during child initialization
    self.assertRaises(AttributeError, lambda: child_1.child_inst_var)
    self.assertEqual(Parent.counter, 0)
    self.assertEqual(Child_1.counter, 0)

    child_2 = Child_Error_2(5)
    self.assertEqual(child_2.child_class_var, 2001)

    # error during child initialization
    self.assertRaises(AttributeError, lambda: child_2.parent_inst_var)
    self.assertEqual(Parent.counter, 0)
    self.assertEqual(Child_2.counter, 0)


def my_normal_func(x, y):
  return x + y


@specifiable
class Wrapper():
  def __init__(self, func=None, cls=None, **kwargs):
    self._func = func
    if cls is not None:
      self._cls = cls(**kwargs)

  def run_func(self, x, y):
    return self._func(x, y)

  def run_func_in_class(self, x, y):
    return self._cls.apply(x, y)


class TestFunctionAsArgument(unittest.TestCase):
  def setUp(self) -> None:
    self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE)

  def tearDown(self) -> None:
    _KNOWN_SPECIFIABLE.clear()
    _KNOWN_SPECIFIABLE.update(self.saved_specifiable)

  def test_normal_function(self):
    w = Wrapper(my_normal_func)

    self.assertEqual(w.run_func(1, 2), 3)

    w_spec = w.to_spec()
    self.assertEqual(
        w_spec,
        Spec(
            type='Wrapper',
            config={'func': Spec(type="my_normal_func", config=None)}))

    w_2 = Specifiable.from_spec(w_spec)
    self.assertEqual(w_2.run_func(2, 3), 5)

  def test_lambda_function(self):
    my_lambda_func = lambda x, y: x - y

    w = Wrapper(my_lambda_func)

    self.assertEqual(w.run_func(3, 2), 1)

    w_spec = w.to_spec()
    self.assertEqual(
        w_spec,
        Spec(
            type='Wrapper',
            config={
                'func': Spec(
                    type=
                    f"<lambda at {os.path.basename(__file__)}:{my_lambda_func.__code__.co_firstlineno}>",  # pylint: disable=line-too-long
                    config=None)
            }))

    w_2 = Specifiable.from_spec(w_spec)
    self.assertEqual(w_2.run_func(5, 3), 2)


class TestClassAsArgument(unittest.TestCase):
  def setUp(self) -> None:
    self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE)

  def tearDown(self) -> None:
    _KNOWN_SPECIFIABLE.clear()
    _KNOWN_SPECIFIABLE.update(self.saved_specifiable)

  def test_normal_class(self):
    class InnerClass():
      def __init__(self, multiplier):
        self._multiplier = multiplier

      def apply(self, x, y):
        return x * y * self._multiplier

    w = Wrapper(cls=InnerClass, multiplier=10)
    self.assertEqual(w.run_func_in_class(2, 3), 60)

    w_spec = w.to_spec()
    self.assertEqual(
        w_spec,
        Spec(
            type='Wrapper',
            config={
                'cls': Spec(type='InnerClass', config=None), 'multiplier': 10
            }))

    w_2 = Specifiable.from_spec(w_spec)
    self.assertEqual(w_2.run_func_in_class(5, 3), 150)


class TestUncommonUsages(unittest.TestCase):
  def test_double_specifiable(self):
    @specifiable
    @specifiable
    class ZZ():
      def __init__(self, a):
        self.a = a

    assert issubclass(ZZ, Specifiable)
    c = ZZ("b")
    c.run_original_init()
    self.assertEqual(c.a, "b")

  def test_unspecifiable(self):
    class YY():
      def __init__(self, x):
        self.x = x
        assert False

    YY = specifiable(YY)
    assert issubclass(YY, Specifiable)
    y = YY(1)
    # __init__ is called (with assertion error raised) when attribute is first
    # accessed
    self.assertRaises(AssertionError, lambda: y.x)

    # unspecifiable YY
    YY.unspecifiable()
    # __init__ is called immediately
    self.assertRaises(AssertionError, YY, 1)
    self.assertFalse(hasattr(YY, 'run_original_init'))
    self.assertFalse(hasattr(YY, 'spec_type'))
    self.assertFalse(hasattr(YY, 'to_spec'))
    self.assertFalse(hasattr(YY, 'from_spec'))
    self.assertFalse(hasattr(YY, 'unspecifiable'))

    # make YY specifiable again
    YY = specifiable(YY)
    assert issubclass(YY, Specifiable)
    y = YY(1)
    self.assertRaises(AssertionError, lambda: y.x)


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