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

"""Unit tests for the Beam Row typing functionality."""

import typing
import unittest
from dataclasses import dataclass

import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.typehints import row_type
from apache_beam.typehints import schemas


class RowTypeTest(unittest.TestCase):
  @staticmethod
  def _check_key_type_and_count(x) -> int:
    key_type = type(x[0])
    if row_type._user_type_is_generated(key_type):
      raise RuntimeError("Type after GBK not preserved, get generated type")
    if not hasattr(key_type, row_type._BEAM_SCHEMA_ID):
      raise RuntimeError("Type after GBK missing Beam schema ID")

    return len(x[1])

  def test_group_by_key_namedtuple(self):
    MyNamedTuple = typing.NamedTuple(
        "MyNamedTuple", [("id", int), ("name", str)])

    beam.coders.typecoders.registry.register_row(MyNamedTuple)

    def generate(num: int):
      for i in range(100):
        yield (MyNamedTuple(i, 'a'), num)

    pipeline = TestPipeline(is_integration_test=False)

    with pipeline as p:
      result = (
          p
          | 'Create' >> beam.Create([i for i in range(10)])
          | 'Generate' >> beam.ParDo(generate).with_output_types(
              tuple[MyNamedTuple, int])
          | 'GBK' >> beam.GroupByKey()
          | 'Count Elements' >> beam.Map(self._check_key_type_and_count))
      assert_that(result, equal_to([10] * 100))

  def test_group_by_key_dataclass(self):
    @dataclass
    class MyDataClass:
      id: int
      name: str

    beam.coders.typecoders.registry.register_row(MyDataClass)

    def generate(num: int):
      for i in range(100):
        yield (MyDataClass(i, 'a'), num)

    pipeline = TestPipeline(is_integration_test=False)

    with pipeline as p:
      result = (
          p
          | 'Create' >> beam.Create([i for i in range(10)])
          | 'Generate' >> beam.ParDo(generate).with_output_types(
              tuple[MyDataClass, int])
          | 'GBK' >> beam.GroupByKey()
          | 'Count Elements' >> beam.Map(self._check_key_type_and_count))
      assert_that(result, equal_to([10] * 100))

  def test_group_by_key_namedtuple_union(self):
    Tuple1 = typing.NamedTuple("Tuple1", [("id", int)])

    Tuple2 = typing.NamedTuple("Tuple2", [("id", int), ("name", str)])

    def generate(num: int):
      for i in range(2):
        yield (Tuple1(i), num)
        yield (Tuple2(i, 'a'), num)

    pipeline = TestPipeline(is_integration_test=False)

    with pipeline as p:
      result = (
          p
          | 'Create' >> beam.Create([i for i in range(2)])
          | 'Generate' >> beam.ParDo(generate).with_output_types(
              tuple[(Tuple1 | Tuple2), int])
          | 'GBK' >> beam.GroupByKey()
          | 'Count' >> beam.Map(lambda x: len(x[1])))
      assert_that(result, equal_to([2] * 4))

  # Union of dataclasses as type hint currently result in FastPrimitiveCoder
  # fails at GBK
  @unittest.skip("https://github.com/apache/beam/issues/22085")
  def test_group_by_key_inherited_dataclass_union(self):
    @dataclass
    class DataClassInt:
      id: int

    @dataclass
    class DataClassStr(DataClassInt):
      name: str

    beam.coders.typecoders.registry.register_row(DataClassInt)
    beam.coders.typecoders.registry.register_row(DataClassStr)

    def generate(num: int):
      for i in range(10):
        yield (DataClassInt(i), num)
        yield (DataClassStr(i, 'a'), num)

    pipeline = TestPipeline(is_integration_test=False)

    with pipeline as p:
      result = (
          p
          | 'Create' >> beam.Create([i for i in range(2)])
          | 'Generate' >> beam.ParDo(generate).with_output_types(
              tuple[(DataClassInt | DataClassStr), int])
          | 'GBK' >> beam.GroupByKey()
          | 'Count Elements' >> beam.Map(self._check_key_type_and_count))
      assert_that(result, equal_to([2] * 4))

  def test_derived_dataclass_schema_id(self):
    @dataclass
    class BaseDataClass:
      id: int

    @dataclass
    class DerivedDataClass(BaseDataClass):
      name: str

    self.assertFalse(hasattr(BaseDataClass, row_type._BEAM_SCHEMA_ID))
    schema_for_base = schemas.schema_from_element_type(BaseDataClass)
    self.assertTrue(hasattr(BaseDataClass, row_type._BEAM_SCHEMA_ID))
    self.assertEqual(
        schema_for_base.id, getattr(BaseDataClass, row_type._BEAM_SCHEMA_ID))

    # Getting the schema for BaseDataClass sets the _beam_schema_id
    schemas.typing_to_runner_api(
        BaseDataClass, schema_registry=schemas.SchemaTypeRegistry())

    # We create a RowTypeConstraint from DerivedDataClass.
    # It should not inherit the _beam_schema_id from BaseDataClass!
    derived_row_type = row_type.RowTypeConstraint.from_user_type(
        DerivedDataClass)
    self.assertIsNone(derived_row_type._schema_id)

    schema_for_derived = schemas.schema_from_element_type(DerivedDataClass)
    self.assertTrue(hasattr(DerivedDataClass, row_type._BEAM_SCHEMA_ID))
    self.assertEqual(
        schema_for_derived.id,
        getattr(DerivedDataClass, row_type._BEAM_SCHEMA_ID))
    self.assertNotEqual(schema_for_derived.id, schema_for_base.id)


if __name__ == '__main__':
  unittest.main()
