/*
 * 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.
 */
package org.apache.beam.fn.harness.data;

import static org.apache.beam.sdk.util.CoderUtils.encodeToByteArray;
import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow;
import static org.hamcrest.Matchers.contains;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import java.util.Collection;
import java.util.UUID;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.Target;
import org.apache.beam.model.fnexecution.v1.BeamFnDataGrpc;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.LengthPrefixCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.fn.data.InboundDataClient;
import org.apache.beam.sdk.fn.data.LogicalEndpoint;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.fn.test.TestExecutors;
import org.apache.beam.sdk.fn.test.TestExecutors.TestExecutorService;
import org.apache.beam.sdk.fn.test.TestStreams;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p13p1.io.grpc.ManagedChannel;
import org.apache.beam.vendor.grpc.v1p13p1.io.grpc.Server;
import org.apache.beam.vendor.grpc.v1p13p1.io.grpc.inprocess.InProcessChannelBuilder;
import org.apache.beam.vendor.grpc.v1p13p1.io.grpc.inprocess.InProcessServerBuilder;
import org.apache.beam.vendor.grpc.v1p13p1.io.grpc.stub.CallStreamObserver;
import org.apache.beam.vendor.grpc.v1p13p1.io.grpc.stub.StreamObserver;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Tests for {@link QueueingBeamFnDataClient}. */
@RunWith(JUnit4.class)
public class QueueingBeamFnDataClientTest {

  private static final Logger LOG = LoggerFactory.getLogger(QueueingBeamFnDataClientTest.class);

  @Rule public TestExecutorService executor = TestExecutors.from(Executors::newCachedThreadPool);

  private static final Coder<WindowedValue<String>> CODER =
      LengthPrefixCoder.of(
          WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE));
  private static final LogicalEndpoint ENDPOINT_A =
      LogicalEndpoint.of(
          "12L",
          Target.newBuilder().setPrimitiveTransformReference("34L").setName("targetA").build());

  private static final LogicalEndpoint ENDPOINT_B =
      LogicalEndpoint.of(
          "56L",
          BeamFnApi.Target.newBuilder()
              .setPrimitiveTransformReference("78L")
              .setName("targetB")
              .build());

  private static final BeamFnApi.Elements ELEMENTS_A_1;
  private static final BeamFnApi.Elements ELEMENTS_A_2;
  private static final BeamFnApi.Elements ELEMENTS_B_1;

  static {
    try {
      ELEMENTS_A_1 =
          BeamFnApi.Elements.newBuilder()
              .addData(
                  BeamFnApi.Elements.Data.newBuilder()
                      .setInstructionReference(ENDPOINT_A.getInstructionId())
                      .setTarget(ENDPOINT_A.getTarget())
                      .setData(
                          ByteString.copyFrom(encodeToByteArray(CODER, valueInGlobalWindow("ABC")))
                              .concat(
                                  ByteString.copyFrom(
                                      encodeToByteArray(CODER, valueInGlobalWindow("DEF"))))))
              .build();
      ELEMENTS_A_2 =
          BeamFnApi.Elements.newBuilder()
              .addData(
                  BeamFnApi.Elements.Data.newBuilder()
                      .setInstructionReference(ENDPOINT_A.getInstructionId())
                      .setTarget(ENDPOINT_A.getTarget())
                      .setData(
                          ByteString.copyFrom(
                              encodeToByteArray(CODER, valueInGlobalWindow("GHI")))))
              .addData(
                  BeamFnApi.Elements.Data.newBuilder()
                      .setInstructionReference(ENDPOINT_A.getInstructionId())
                      .setTarget(ENDPOINT_A.getTarget()))
              .build();
      ELEMENTS_B_1 =
          BeamFnApi.Elements.newBuilder()
              .addData(
                  BeamFnApi.Elements.Data.newBuilder()
                      .setInstructionReference(ENDPOINT_B.getInstructionId())
                      .setTarget(ENDPOINT_B.getTarget())
                      .setData(
                          ByteString.copyFrom(encodeToByteArray(CODER, valueInGlobalWindow("JKL")))
                              .concat(
                                  ByteString.copyFrom(
                                      encodeToByteArray(CODER, valueInGlobalWindow("MNO"))))))
              .addData(
                  BeamFnApi.Elements.Data.newBuilder()
                      .setInstructionReference(ENDPOINT_B.getInstructionId())
                      .setTarget(ENDPOINT_B.getTarget()))
              .build();
    } catch (Exception e) {
      throw new ExceptionInInitializerError(e);
    }
  }

  @Test(timeout = 10000)
  public void testBasicInboundConsumerBehaviour() throws Exception {
    CountDownLatch waitForClientToConnect = new CountDownLatch(1);
    CountDownLatch receiveAllValuesA = new CountDownLatch(3);
    CountDownLatch receiveAllValuesB = new CountDownLatch(2);
    Collection<WindowedValue<String>> inboundValuesA = new ConcurrentLinkedQueue<>();
    Collection<WindowedValue<String>> inboundValuesB = new ConcurrentLinkedQueue<>();
    Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.Elements>> outboundServerObserver =
        new AtomicReference<>();
    CallStreamObserver<BeamFnApi.Elements> inboundServerObserver =
        TestStreams.withOnNext(inboundServerValues::add).build();

    Endpoints.ApiServiceDescriptor apiServiceDescriptor =
        Endpoints.ApiServiceDescriptor.newBuilder()
            .setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString())
            .build();
    Server server =
        InProcessServerBuilder.forName(apiServiceDescriptor.getUrl())
            .addService(
                new BeamFnDataGrpc.BeamFnDataImplBase() {
                  @Override
                  public StreamObserver<BeamFnApi.Elements> data(
                      StreamObserver<BeamFnApi.Elements> outboundObserver) {
                    outboundServerObserver.set(outboundObserver);
                    waitForClientToConnect.countDown();
                    return inboundServerObserver;
                  }
                })
            .build();
    server.start();
    try {
      ManagedChannel channel =
          InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();

      BeamFnDataGrpcClient clientFactory =
          new BeamFnDataGrpcClient(
              PipelineOptionsFactory.create(),
              (Endpoints.ApiServiceDescriptor descriptor) -> channel,
              OutboundObserverFactory.trivial());
      QueueingBeamFnDataClient queueingClient = new QueueingBeamFnDataClient(clientFactory);

      InboundDataClient readFutureA =
          queueingClient.receive(
              apiServiceDescriptor,
              ENDPOINT_A,
              CODER,
              (WindowedValue<String> wv) -> {
                inboundValuesA.add(wv);
                receiveAllValuesA.countDown();
              });

      waitForClientToConnect.await();

      Future<?> sendElementsFuture =
          executor.submit(
              () -> {
                outboundServerObserver.get().onNext(ELEMENTS_A_1);
                // Purposefully transmit some data before the consumer for B is bound showing that
                // data is not lost
                outboundServerObserver.get().onNext(ELEMENTS_B_1);
              });

      // This can be compeleted before we get values?
      InboundDataClient readFutureB =
          queueingClient.receive(
              apiServiceDescriptor,
              ENDPOINT_B,
              CODER,
              (WindowedValue<String> wv) -> {
                inboundValuesB.add(wv);
                receiveAllValuesB.countDown();
              });

      Future<?> drainElementsFuture =
          executor.submit(
              () -> {
                try {
                  queueingClient.drainAndBlock();
                } catch (Exception e) {
                  LOG.error("Failed ", e);
                  fail();
                }
              });

      receiveAllValuesB.await();
      assertThat(inboundValuesB, contains(valueInGlobalWindow("JKL"), valueInGlobalWindow("MNO")));

      outboundServerObserver.get().onNext(ELEMENTS_A_2);

      receiveAllValuesA.await(); // Wait for A's values to be available
      assertThat(
          inboundValuesA,
          contains(
              valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF"), valueInGlobalWindow("GHI")));

      // Wait for these threads to terminate
      sendElementsFuture.get();
      drainElementsFuture.get();
    } finally {
      server.shutdownNow();
    }
  }

  @Test(timeout = 10000)
  public void testBundleProcessorThrowsExecutionExceptionWhenUserCodeThrows() throws Exception {
    CountDownLatch waitForClientToConnect = new CountDownLatch(1);
    //Collection<WindowedValue<String>> inboundValuesA = new ConcurrentLinkedQueue<>();
    Collection<WindowedValue<String>> inboundValuesB = new ConcurrentLinkedQueue<>();
    Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.Elements>> outboundServerObserver =
        new AtomicReference<>();
    CallStreamObserver<BeamFnApi.Elements> inboundServerObserver =
        TestStreams.withOnNext(inboundServerValues::add).build();

    Endpoints.ApiServiceDescriptor apiServiceDescriptor =
        Endpoints.ApiServiceDescriptor.newBuilder()
            .setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString())
            .build();
    Server server =
        InProcessServerBuilder.forName(apiServiceDescriptor.getUrl())
            .addService(
                new BeamFnDataGrpc.BeamFnDataImplBase() {
                  @Override
                  public StreamObserver<BeamFnApi.Elements> data(
                      StreamObserver<BeamFnApi.Elements> outboundObserver) {
                    outboundServerObserver.set(outboundObserver);
                    waitForClientToConnect.countDown();
                    return inboundServerObserver;
                  }
                })
            .build();
    server.start();
    try {
      ManagedChannel channel =
          InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();

      BeamFnDataGrpcClient clientFactory =
          new BeamFnDataGrpcClient(
              PipelineOptionsFactory.create(),
              (Endpoints.ApiServiceDescriptor descriptor) -> channel,
              OutboundObserverFactory.trivial());
      QueueingBeamFnDataClient queueingClient = new QueueingBeamFnDataClient(clientFactory);

      InboundDataClient readFutureA =
          queueingClient.receive(
              apiServiceDescriptor,
              ENDPOINT_A,
              CODER,
              (WindowedValue<String> wv) -> {
                throw new RuntimeException("Intentionally fail!"); // Error injected here.
              });

      waitForClientToConnect.await();

      Future<?> sendElementsFuture =
          executor.submit(
              () -> {
                outboundServerObserver.get().onNext(ELEMENTS_A_1);
                // Purposefully transmit some data before the consumer for B is bound showing that
                // data is not lost
                outboundServerObserver.get().onNext(ELEMENTS_B_1);
              });

      InboundDataClient readFutureB =
          queueingClient.receive(
              apiServiceDescriptor,
              ENDPOINT_B,
              CODER,
              (WindowedValue<String> wv) -> {
                inboundValuesB.add(wv);
              });

      Future<?> drainElementsFuture =
          executor.submit(
              () -> {
                boolean intentionallyFailed = false;
                try {
                  queueingClient.drainAndBlock();
                } catch (RuntimeException e) {
                  intentionallyFailed = true;
                } catch (Exception e) {
                  LOG.error("Unintentional failure", e);
                  fail();
                }
                assertTrue(intentionallyFailed);
              });

      // Fail all InboundObservers if any of the downstream consumers fail.
      // This allows the ProcessBundlerHandler to unblock everything and fail properly.
      boolean intentionallyFailedA = false;
      try {
        readFutureA.awaitCompletion();
      } catch (ExecutionException e) {
        if (e.getCause() instanceof RuntimeException) {
          intentionallyFailedA = true;
        }
      }
      assertTrue(intentionallyFailedA);

      boolean intentionallyFailedB = false;
      try {
        readFutureB.awaitCompletion();
      } catch (ExecutionException e) {
        if (e.getCause() instanceof RuntimeException) {
          intentionallyFailedB = true;
        }
      } catch (Exception e) {
        intentionallyFailedB = true;
      }
      assertTrue(intentionallyFailedB);

      // Wait for these threads to terminate
      sendElementsFuture.get();
      drainElementsFuture.get();
    } finally {
      server.shutdownNow();
    }
  }
}
