#
# 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 for the Java external transforms."""

import argparse
import logging
import subprocess
import sys
from typing import Optional

import grpc
from mock import patch

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
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.transforms.external import ImplicitSchemaPayloadBuilder
from apache_beam.utils import subprocess_server

# Protect against environments where apitools library is not available.
# pylint: disable=wrong-import-order, wrong-import-position
try:
  from apache_beam.runners.dataflow.internal import apiclient as _apiclient
except ImportError:
  apiclient = None
else:
  apiclient = _apiclient
# pylint: enable=wrong-import-order, wrong-import-position


class JavaExternalTransformTest(object):

  # This will be overwritten if set via a flag.
  expansion_service_jar: Optional[str] = None
  expansion_service_port: Optional[int] = None

  class _RunWithExpansion(object):
    def __init__(self):
      self._server = None

    def __enter__(self):
      if not (JavaExternalTransformTest.expansion_service_jar or
              JavaExternalTransformTest.expansion_service_port):
        raise RuntimeError('No expansion service jar or port provided.')

      JavaExternalTransformTest.expansion_service_port = (
          JavaExternalTransformTest.expansion_service_port or 8091)

      jar = JavaExternalTransformTest.expansion_service_jar
      port = JavaExternalTransformTest.expansion_service_port

      # Start the java server and wait for it to be ready.
      if jar:
        java = subprocess_server.JavaHelper.get_java()
        self._server = subprocess.Popen([java, '-jar', jar, str(port)])

      address = 'localhost:%s' % str(port)

      with grpc.insecure_channel(address) as channel:
        grpc.channel_ready_future(channel).result()

    def __exit__(self, type, value, traceback):
      if self._server:
        self._server.kill()
        self._server = None

  @staticmethod
  def test_java_expansion_dataflow():
    if apiclient is None:
      return

    # This test does not actually running the pipeline in Dataflow. It just
    # tests the translation to a Dataflow job request.

    with patch.object(apiclient.DataflowApplicationClient,
                      'create_job') as mock_create_job:
      with JavaExternalTransformTest._RunWithExpansion():
        pipeline_options = PipelineOptions([
            '--runner=DataflowRunner',
            '--project=dummyproject',
            '--region=some-region1',
            '--experiments=beam_fn_api',
            '--temp_location=gs://dummybucket/'
        ])

        # Run a simple count-filtered-letters pipeline.
        JavaExternalTransformTest.run_pipeline(
            pipeline_options,
            JavaExternalTransformTest.expansion_service_port,
            False)

        mock_args = mock_create_job.call_args_list
        assert mock_args
        args, kwargs = mock_args[0]
        job = args[0]
        job_str = '%s' % job
        assert 'beam:transforms:xlang:filter_less_than_eq' in job_str

  @staticmethod
  def run_pipeline_with_expansion_service(pipeline_options):
    with JavaExternalTransformTest._RunWithExpansion():
      # Run a simple count-filtered-letters pipeline.
      JavaExternalTransformTest.run_pipeline(
          pipeline_options,
          JavaExternalTransformTest.expansion_service_port,
          True)

  @staticmethod
  def run_pipeline(pipeline_options, expansion_service, wait_until_finish=True):
    # The actual definitions of these transforms is in
    # org.apache.beam.sdk.util.construction.TestExpansionService.
    TEST_COUNT_URN = "beam:transforms:xlang:count"
    TEST_FILTER_URN = "beam:transforms:xlang:filter_less_than_eq"

    # Run a simple count-filtered-letters pipeline.
    p = TestPipeline(options=pipeline_options)

    if isinstance(expansion_service, int):
      # Only the port was specified.
      expansion_service = 'localhost:%s' % str(expansion_service)

    res = (
        p
        | beam.Create(list('aaabccxyyzzz'))
        | beam.Map(str)
        | beam.ExternalTransform(
            TEST_FILTER_URN,
            ImplicitSchemaPayloadBuilder({'data': 'middle'}),
            expansion_service)
        | beam.ExternalTransform(TEST_COUNT_URN, None, expansion_service)
        # Map(lambda) produces a label formatted like this, but it cannot be
        # changed without breaking update compat. Here, we pin to the transform
        # name used in the 2.68 release to avoid breaking changes when the line
        # number changes. Context: https://github.com/apache/beam/pull/36381
        | "Map(<lambda at external_java.py:148>)" >>
        beam.Map(lambda kv: '%s: %s' % kv))

    assert_that(res, equal_to(['a: 3', 'b: 1', 'c: 2']))

    result = p.run()
    if wait_until_finish:
      result.wait_until_finish()


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  parser = argparse.ArgumentParser()
  parser.add_argument('--expansion_service_jar')
  parser.add_argument('--expansion_service_port')
  parser.add_argument('--expansion_service_target')
  parser.add_argument('--expansion_service_target_appendix')
  known_args, pipeline_args = parser.parse_known_args(sys.argv)

  if known_args.expansion_service_jar:
    JavaExternalTransformTest.expansion_service_jar = (
        known_args.expansion_service_jar)
    JavaExternalTransformTest.expansion_service_port = int(
        known_args.expansion_service_port)
    pipeline_options = PipelineOptions(pipeline_args)
    JavaExternalTransformTest.run_pipeline_with_expansion_service(
        pipeline_options)
  elif known_args.expansion_service_target:
    pipeline_options = PipelineOptions(pipeline_args)
    JavaExternalTransformTest.run_pipeline(
        pipeline_options,
        beam.transforms.external.BeamJarExpansionService(
            known_args.expansion_service_target,
            gradle_appendix=known_args.expansion_service_target_appendix))
  else:
    raise RuntimeError(
        "--expansion_service_jar or --expansion_service_target "
        "should be provided.")