/*
 * 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.runners.dataflow.worker.windmill.client.grpc;

import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableListMultimap.toImmutableListMultimap;

import com.google.auto.value.AutoBuilder;
import java.io.PrintWriter;
import java.util.Collection;
import java.util.List;
import java.util.Set;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.runners.dataflow.worker.status.StatusDataProvider;
import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillMetadataServiceV1Alpha1Grpc.CloudWindmillMetadataServiceV1Alpha1Stub;
import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationHeartbeatResponse;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetWorkRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillConnection;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
import org.apache.beam.runners.dataflow.worker.windmill.client.AbstractWindmillStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkerMetadataStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.commits.WorkCommitter;
import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.GetDataClient;
import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.StreamObserverFactory;
import org.apache.beam.runners.dataflow.worker.windmill.work.WorkItemReceiver;
import org.apache.beam.runners.dataflow.worker.windmill.work.WorkItemScheduler;
import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.util.BackOff;
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.AbstractStub;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Suppliers;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.joda.time.Duration;
import org.joda.time.Instant;

/**
 * Creates gRPC streaming connections to Windmill Service. Maintains a set of all currently opened
 * RPC streams for health check/heartbeat requests to keep the streams alive.
 */
@ThreadSafe
@Internal
public class GrpcWindmillStreamFactory implements StatusDataProvider {
  private static final long DEFAULT_STREAM_RPC_DEADLINE_SECONDS = 300;
  private static final java.time.Duration
      DEFAULT_DIRECT_STREAMING_RPC_PHYSICAL_STREAM_HALF_CLOSE_AFTER =
          java.time.Duration.ofMinutes(3);
  private static final Duration MIN_BACKOFF = Duration.millis(1);
  private static final Duration DEFAULT_MAX_BACKOFF = Duration.standardSeconds(30);
  private static final int DEFAULT_LOG_EVERY_N_STREAM_FAILURES = 1;
  private static final int DEFAULT_STREAMING_RPC_BATCH_LIMIT = Integer.MAX_VALUE;
  private static final int DEFAULT_WINDMILL_MESSAGES_BETWEEN_IS_READY_CHECKS = 1;
  private static final int NO_HEALTH_CHECKS = -1;
  private static final String NO_BACKEND_WORKER_TOKEN = "";
  private static final String DISPATCHER_DEBUG_NAME = "Dispatcher";

  private final JobHeader jobHeader;
  private final int logEveryNStreamFailures;
  private final int streamingRpcBatchLimit;
  private final int windmillMessagesBetweenIsReadyChecks;
  private final Supplier<BackOff> grpcBackOff;
  private final Set<AbstractWindmillStream<?, ?>> streamRegistry;
  private final AtomicLong streamIdGenerator;
  // If true, then active work refreshes will be sent as KeyedGetDataRequests. Otherwise, use the
  // newer ComputationHeartbeatRequests.
  private final boolean sendKeyedGetDataRequests;
  private final boolean requestBatchedGetWorkResponse;
  private final Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses;
  private final java.time.Duration directStreamingRpcPhysicalStreamHalfCloseAfter;
  private final Supplier<ScheduledExecutorService> executorServiceSupplier;
  private final java.time.Duration commitWorkStreamRetryTimeout;

  private GrpcWindmillStreamFactory(
      JobHeader jobHeader,
      int logEveryNStreamFailures,
      int streamingRpcBatchLimit,
      int windmillMessagesBetweenIsReadyChecks,
      boolean sendKeyedGetDataRequests,
      boolean requestBatchedGetWorkResponse,
      Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses,
      Supplier<Duration> maxBackOffSupplier,
      java.time.Duration directStreamingRpcPhysicalStreamHalfCloseAfter,
      java.time.Duration commitWorkStreamRetryTimeout,
      Supplier<ScheduledExecutorService> executorServiceSupplier) {
    this.jobHeader = jobHeader;
    this.logEveryNStreamFailures = logEveryNStreamFailures;
    this.streamingRpcBatchLimit = streamingRpcBatchLimit;
    this.windmillMessagesBetweenIsReadyChecks = windmillMessagesBetweenIsReadyChecks;
    // Configure backoff to retry calls forever, with a maximum sane retry interval.
    Supplier<FluentBackoff> backoffConfig =
        Suppliers.memoize(
            () ->
                FluentBackoff.DEFAULT
                    .withInitialBackoff(MIN_BACKOFF)
                    .withMaxBackoff(maxBackOffSupplier.get()));
    this.grpcBackOff = () -> backoffConfig.get().backoff();
    this.streamRegistry = ConcurrentHashMap.newKeySet();
    this.sendKeyedGetDataRequests = sendKeyedGetDataRequests;
    this.requestBatchedGetWorkResponse = requestBatchedGetWorkResponse;
    this.processHeartbeatResponses = processHeartbeatResponses;
    this.streamIdGenerator = new AtomicLong();
    this.directStreamingRpcPhysicalStreamHalfCloseAfter =
        directStreamingRpcPhysicalStreamHalfCloseAfter;
    this.executorServiceSupplier = executorServiceSupplier;
    this.commitWorkStreamRetryTimeout = commitWorkStreamRetryTimeout;
  }

  /** @implNote Used for {@link AutoBuilder} {@link Builder} class, do not call directly. */
  static GrpcWindmillStreamFactory create(
      JobHeader jobHeader,
      int logEveryNStreamFailures,
      int streamingRpcBatchLimit,
      int windmillMessagesBetweenIsReadyChecks,
      boolean sendKeyedGetDataRequests,
      boolean requestBatchedGetWorkResponse,
      Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses,
      Supplier<Duration> maxBackOffSupplier,
      int healthCheckIntervalMillis,
      java.time.Duration directStreamingRpcPhysicalStreamHalfCloseAfter,
      java.time.Duration commitWorkStreamRetryTimeout,
      Supplier<ScheduledExecutorService> scheduledExecutorServiceSupplier) {
    GrpcWindmillStreamFactory streamFactory =
        new GrpcWindmillStreamFactory(
            jobHeader,
            logEveryNStreamFailures,
            streamingRpcBatchLimit,
            windmillMessagesBetweenIsReadyChecks,
            sendKeyedGetDataRequests,
            requestBatchedGetWorkResponse,
            processHeartbeatResponses,
            maxBackOffSupplier,
            directStreamingRpcPhysicalStreamHalfCloseAfter,
            commitWorkStreamRetryTimeout,
            scheduledExecutorServiceSupplier);

    if (healthCheckIntervalMillis >= 0) {
      // Health checks are run on background daemon thread, which will only be cleaned up on JVM
      // shutdown.
      new Timer("WindmillHealthCheckTimer")
          .schedule(
              new TimerTask() {
                @Override
                public void run() {
                  Instant reportThreshold =
                      Instant.now().minus(Duration.millis(healthCheckIntervalMillis));
                  for (AbstractWindmillStream<?, ?> stream : streamFactory.streamRegistry) {
                    stream.maybeScheduleHealthCheck(reportThreshold);
                  }
                }
              },
              0,
              healthCheckIntervalMillis);
    }

    return streamFactory;
  }

  /**
   * Returns a new {@link Builder} for {@link GrpcWindmillStreamFactory} with default values set for
   * the given {@link JobHeader}.
   */
  @SuppressWarnings("nullness")
  public static GrpcWindmillStreamFactory.Builder of(JobHeader jobHeader) {
    return new AutoBuilder_GrpcWindmillStreamFactory_Builder()
        .setJobHeader(jobHeader)
        .setWindmillMessagesBetweenIsReadyChecks(DEFAULT_WINDMILL_MESSAGES_BETWEEN_IS_READY_CHECKS)
        .setMaxBackOffSupplier(() -> DEFAULT_MAX_BACKOFF)
        .setLogEveryNStreamFailures(DEFAULT_LOG_EVERY_N_STREAM_FAILURES)
        .setStreamingRpcBatchLimit(DEFAULT_STREAMING_RPC_BATCH_LIMIT)
        .setHealthCheckIntervalMillis(NO_HEALTH_CHECKS)
        .setSendKeyedGetDataRequests(true)
        .setRequestBatchedGetWorkResponse(false)
        .setProcessHeartbeatResponses(ignored -> {})
        .setDirectStreamingRpcPhysicalStreamHalfCloseAfter(
            DEFAULT_DIRECT_STREAMING_RPC_PHYSICAL_STREAM_HALF_CLOSE_AFTER)
        .setCommitWorkStreamRetryTimeout(java.time.Duration.ZERO)
        .setScheduledExecutorServiceSupplier(() -> null);
  }

  private static <T extends AbstractStub<T>> T withDefaultDeadline(T stub) {
    // Deadlines are absolute points in time, so generate a new one everytime this function is
    // called.
    return stub.withDeadlineAfter(DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);
  }

  private static void printSummaryHtmlForWorker(
      String workerToken, Collection<AbstractWindmillStream<?, ?>> streams, PrintWriter writer) {
    writer.write(
        "<strong>" + (workerToken.isEmpty() ? DISPATCHER_DEBUG_NAME : workerToken) + "</strong>");
    writer.write("<br>");
    streams.forEach(
        stream -> {
          stream.appendSummaryHtml(writer);
          writer.write("<br>");
        });
    writer.write("<br>");
  }

  private ScheduledExecutorService executorForDispatchedStreams(String debugStreamTypeName) {
    ScheduledExecutorService result = executorServiceSupplier.get();
    if (result != null) {
      return result;
    }
    return Executors.newSingleThreadScheduledExecutor(
        new ThreadFactoryBuilder()
            .setDaemon(true)
            .setNameFormat(String.format("%s-WindmillStream-thread", debugStreamTypeName))
            .build());
  }

  private ScheduledExecutorService executorForDirectStreams(
      String backendWorkerToken, String debugStreamTypeName) {
    ScheduledExecutorService supplierResult = executorServiceSupplier.get();
    if (supplierResult != null) {
      return supplierResult;
    }
    ScheduledThreadPoolExecutor result =
        new ScheduledThreadPoolExecutor(
            0,
            new ThreadFactoryBuilder()
                .setDaemon(true)
                .setNameFormat(
                    String.join(
                        "-",
                        debugStreamTypeName,
                        backendWorkerToken.substring(0, Math.min(10, backendWorkerToken.length())),
                        "WindmillStream",
                        "%d"))
                .build());
    result.setKeepAliveTime(1, TimeUnit.MINUTES);
    return result;
  }

  public GetWorkStream createGetWorkStream(
      CloudWindmillServiceV1Alpha1Stub stub,
      GetWorkRequest request,
      WorkItemReceiver processWorkItem) {
    return GrpcGetWorkStream.create(
        NO_BACKEND_WORKER_TOKEN,
        responseObserver -> withDefaultDeadline(stub).getWorkStream(responseObserver),
        request,
        grpcBackOff.get(),
        newStreamObserverFactory(),
        streamRegistry,
        logEveryNStreamFailures,
        requestBatchedGetWorkResponse,
        processWorkItem,
        java.time.Duration.ZERO,
        executorForDispatchedStreams("GetWork"));
  }

  public GetWorkStream createDirectGetWorkStream(
      WindmillConnection connection,
      GetWorkRequest request,
      HeartbeatSender heartbeatSender,
      GetDataClient getDataClient,
      WorkCommitter workCommitter,
      WorkItemScheduler workItemScheduler) {
    return GrpcDirectGetWorkStream.create(
        connection.backendWorkerToken(),
        responseObserver ->
            withDefaultDeadline(connection.currentStub()).getWorkStream(responseObserver),
        request,
        grpcBackOff.get(),
        newStreamObserverFactory(),
        streamRegistry,
        logEveryNStreamFailures,
        requestBatchedGetWorkResponse,
        heartbeatSender,
        getDataClient,
        workCommitter,
        workItemScheduler,
        directStreamingRpcPhysicalStreamHalfCloseAfter,
        executorForDirectStreams(connection.backendWorkerToken(), "GetWork"));
  }

  public GetDataStream createGetDataStream(CloudWindmillServiceV1Alpha1Stub stub) {
    return GrpcGetDataStream.create(
        NO_BACKEND_WORKER_TOKEN,
        responseObserver -> withDefaultDeadline(stub).getDataStream(responseObserver),
        grpcBackOff.get(),
        newStreamObserverFactory(),
        streamRegistry,
        logEveryNStreamFailures,
        jobHeader,
        streamIdGenerator,
        streamingRpcBatchLimit,
        sendKeyedGetDataRequests,
        processHeartbeatResponses,
        java.time.Duration.ZERO,
        executorForDispatchedStreams("GetWorkerMetadata"));
  }

  public GetDataStream createDirectGetDataStream(WindmillConnection connection) {
    return GrpcGetDataStream.create(
        connection.backendWorkerToken(),
        responseObserver ->
            withDefaultDeadline(connection.currentStub()).getDataStream(responseObserver),
        grpcBackOff.get(),
        newStreamObserverFactory(),
        streamRegistry,
        logEveryNStreamFailures,
        jobHeader,
        streamIdGenerator,
        streamingRpcBatchLimit,
        sendKeyedGetDataRequests,
        processHeartbeatResponses,
        directStreamingRpcPhysicalStreamHalfCloseAfter,
        executorForDirectStreams(connection.backendWorkerToken(), "GetData"));
  }

  public CommitWorkStream createCommitWorkStream(CloudWindmillServiceV1Alpha1Stub stub) {
    return GrpcCommitWorkStream.create(
        NO_BACKEND_WORKER_TOKEN,
        responseObserver -> withDefaultDeadline(stub).commitWorkStream(responseObserver),
        grpcBackOff.get(),
        newStreamObserverFactory(),
        streamRegistry,
        logEveryNStreamFailures,
        jobHeader,
        streamIdGenerator,
        streamingRpcBatchLimit,
        java.time.Duration.ZERO,
        commitWorkStreamRetryTimeout,
        executorForDispatchedStreams("CommitWork"));
  }

  public CommitWorkStream createDirectCommitWorkStream(WindmillConnection connection) {
    return GrpcCommitWorkStream.create(
        connection.backendWorkerToken(),
        responseObserver ->
            withDefaultDeadline(connection.currentStub()).commitWorkStream(responseObserver),
        grpcBackOff.get(),
        newStreamObserverFactory(),
        streamRegistry,
        logEveryNStreamFailures,
        jobHeader,
        streamIdGenerator,
        streamingRpcBatchLimit,
        directStreamingRpcPhysicalStreamHalfCloseAfter,
        java.time.Duration.ZERO,
        executorForDirectStreams(connection.backendWorkerToken(), "CommitWork"));
  }

  public GetWorkerMetadataStream createGetWorkerMetadataStream(
      Supplier<CloudWindmillMetadataServiceV1Alpha1Stub> stub,
      Consumer<WindmillEndpoints> onNewWindmillEndpoints) {
    return GrpcGetWorkerMetadataStream.create(
        responseObserver -> withDefaultDeadline(stub.get()).getWorkerMetadata(responseObserver),
        grpcBackOff.get(),
        newStreamObserverFactory(),
        streamRegistry,
        logEveryNStreamFailures,
        jobHeader,
        onNewWindmillEndpoints,
        directStreamingRpcPhysicalStreamHalfCloseAfter,
        executorForDispatchedStreams("GetWorkerMetadataStream"));
  }

  private StreamObserverFactory newStreamObserverFactory() {
    return StreamObserverFactory.direct(
        DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, windmillMessagesBetweenIsReadyChecks);
  }

  @Override
  public void appendSummaryHtml(PrintWriter writer) {
    writer.write("Active Streams:<br>");
    streamRegistry.stream()
        .collect(
            toImmutableListMultimap(
                AbstractWindmillStream::backendWorkerToken, Function.identity()))
        .asMap()
        .forEach((workerToken, streams) -> printSummaryHtmlForWorker(workerToken, streams, writer));
  }

  @VisibleForTesting
  final ImmutableSet<AbstractWindmillStream<?, ?>> streamRegistry() {
    return ImmutableSet.copyOf(streamRegistry);
  }

  @Internal
  @AutoBuilder(callMethod = "create")
  public interface Builder {
    Builder setJobHeader(JobHeader jobHeader);

    Builder setLogEveryNStreamFailures(int logEveryNStreamFailures);

    Builder setStreamingRpcBatchLimit(int streamingRpcBatchLimit);

    Builder setWindmillMessagesBetweenIsReadyChecks(int windmillMessagesBetweenIsReadyChecks);

    Builder setMaxBackOffSupplier(Supplier<Duration> maxBackOff);

    Builder setSendKeyedGetDataRequests(boolean sendKeyedGetDataRequests);

    Builder setProcessHeartbeatResponses(
        Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses);

    Builder setHealthCheckIntervalMillis(int healthCheckIntervalMillis);

    Builder setRequestBatchedGetWorkResponse(boolean enabled);

    Builder setDirectStreamingRpcPhysicalStreamHalfCloseAfter(java.time.Duration timeout);

    Builder setCommitWorkStreamRetryTimeout(java.time.Duration timeout);

    Builder setScheduledExecutorServiceSupplier(
        Supplier<ScheduledExecutorService> scheduledExecutorServiceSupplier);

    GrpcWindmillStreamFactory build();
  }
}
