/*
 * 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.spark.util;

import io.opentelemetry.context.Context;
import java.io.Serializable;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.spark.stateful.SparkTimerInternals;
import org.apache.beam.runners.spark.translation.AbstractInOutIterator;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.ValueKind;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext$;
import org.apache.spark.streaming.State;
import org.apache.spark.streaming.StateSpec;
import org.apache.spark.streaming.Time;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.streaming.api.java.JavaPairDStream;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
import scala.Option;
import scala.Tuple2;
import scala.collection.JavaConverters;
import scala.runtime.AbstractFunction3;

/**
 * Utility class for handling timers in the Spark runner. Provides functionality for managing timer
 * operations in stateful processing.
 */
public class TimerUtils {

  public static final byte[] EMPTY_BYTE_ARRAY = new byte[0];

  /**
   * A serializable version of the AbstractFunction3 class from Scala. Used for stateful operations
   * in Spark Streaming.
   *
   * @param <T1> Type of the first parameter
   * @param <T2> Type of the second parameter
   * @param <T3> Type of the third parameter
   * @param <ReturnT> Return type of the function
   */
  abstract static class SerializableFunction3<T1, T2, T3, ReturnT>
      extends AbstractFunction3<T1, T2, T3, ReturnT> implements Serializable {}

  /**
   * A marker class used to identify timer keys and values in Spark transformations. Helps
   * distinguish between regular data flow elements and timer events.
   */
  public static class TimerMarker implements Serializable {
    private static final long serialVersionUID = 1L;

    @Override
    public String toString() {
      return "TIMER_MARKER";
    }

    @Override
    public boolean equals(@Nullable Object o) {
      return o instanceof TimerMarker;
    }

    @Override
    public int hashCode() {
      return 1; // All instances are equal, so they should have the same hash code
    }
  }

  /** Constant marker used to identify timer values in transformations. */
  public static final TimerMarker TIMER_MARKER = new TimerMarker();

  private static class WindowedValueForTimerMarker<T> implements Serializable, WindowedValue<T> {
    private final T value;

    public WindowedValueForTimerMarker(T value) {
      this.value = value;
    }

    @Override
    public PaneInfo getPaneInfo() {
      return PaneInfo.NO_FIRING;
    }

    @Override
    public @Nullable String getRecordId() {
      return null;
    }

    @Override
    public @Nullable Context getOpenTelemetryContext() {
      return null;
    }

    @Override
    public CausedByDrain causedByDrain() {
      return CausedByDrain.NORMAL;
    }

    @Override
    public ValueKind getValueKind() {
      return ValueKind.INSERT;
    }

    @Override
    public @Nullable Long getRecordOffset() {
      return null;
    }

    @Override
    public Iterable<WindowedValue<T>> explodeWindows() {
      return Collections.emptyList();
    }

    @Override
    public T getValue() {
      return value;
    }

    @Override
    public <NewT> WindowedValue<NewT> withValue(NewT newValue) {
      return new WindowedValueForTimerMarker<>(newValue);
    }

    @Override
    public Instant getTimestamp() {
      return BoundedWindow.TIMESTAMP_MIN_VALUE;
    }

    @Override
    public Collection<? extends BoundedWindow> getWindows() {
      return Collections.emptyList();
    }

    @Override
    public boolean equals(@Nullable Object o) {
      if (o instanceof WindowedValueForTimerMarker) {
        WindowedValueForTimerMarker<?> that = (WindowedValueForTimerMarker<?>) o;
        return Objects.equals(this.getValue(), that.getValue());
      } else {
        return super.equals(o);
      }
    }

    @Override
    public int hashCode() {
      return Objects.hash(getValue());
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(getClass())
          .add("value", getValue())
          .add("pane", getPaneInfo())
          .toString();
    }
  }

  /**
   * Fires all expired timers using the provided iterator.
   *
   * <p>Gets expired timers from {@link SparkTimerInternals} and processes each one. The {@link
   * AbstractInOutIterator#fireTimer} method automatically deletes each timer after processing.
   *
   * @param sparkTimerInternals Source of timer data
   * @param windowingStrategy Used to determine which timers are expired
   * @param abstractInOutIterator Iterator that processes and then removes the timers
   * @param <W> Window type
   */
  public static <W extends BoundedWindow> void triggerExpiredTimers(
      SparkTimerInternals sparkTimerInternals,
      WindowingStrategy<?, W> windowingStrategy,
      AbstractInOutIterator<?, ?, ?> abstractInOutIterator) {
    final Collection<TimerInternals.TimerData> expiredTimers =
        getExpiredTimers(sparkTimerInternals, windowingStrategy);

    if (!expiredTimers.isEmpty()) {
      expiredTimers.forEach(abstractInOutIterator::fireTimer);
    }
  }

  public static <W extends BoundedWindow> void dropExpiredTimers(
      SparkTimerInternals sparkTimerInternals, WindowingStrategy<?, W> windowingStrategy) {
    final Collection<TimerInternals.TimerData> expiredTimers =
        getExpiredTimers(sparkTimerInternals, windowingStrategy);

    // Remove the expired timer from the timerInternals structure
    if (!expiredTimers.isEmpty()) {
      expiredTimers.forEach(sparkTimerInternals::deleteTimer);
    }
  }

  private static <W extends BoundedWindow> Collection<TimerInternals.TimerData> getExpiredTimers(
      SparkTimerInternals sparkTimerInternals, WindowingStrategy<?, W> windowingStrategy) {
    return sparkTimerInternals.getTimers().stream()
        .filter(
            timer ->
                timer
                    .getTimestamp()
                    .plus(windowingStrategy.getAllowedLateness())
                    .isBefore(
                        timer.getDomain().equals(TimeDomain.PROCESSING_TIME)
                            ? sparkTimerInternals.currentProcessingTime()
                            : sparkTimerInternals.currentInputWatermarkTime()))
        .collect(Collectors.toList());
  }

  /**
   * Converts a standard DStream into a periodic DStream that ensures all keys are processed in
   * every micro-batch, even if they don't receive new data.
   *
   * <p>This method addresses a fundamental challenge in stateful processing with Spark Streaming:
   * ensuring that timer events for all keys are processed in every batch, regardless of whether
   * those keys receive new data. Without this mechanism, timers associated with inactive keys would
   * not fire at the appropriate times.
   *
   * <p>Implementation details:
   *
   * <ol>
   *   <li>First, it extracts all unique keys from the input DStream and marks them with a special
   *       timer key marker.
   *   <li>Then, it uses Spark's mapWithState operation to maintain a persistent set of all keys
   *       that have ever been seen.
   *   <li>For each batch, it takes a snapshot of this state (all known keys).
   *   <li>It then transforms the original DStream by combining it with synthetic events for all
   *       known keys.
   *   <li>These synthetic events use the special {@link WindowedValueForTimerMarker} class with the
   *       {@link TimerUtils#TIMER_MARKER} to allow downstream operations to distinguish them from
   *       regular data events.
   * </ol>
   *
   * <p>When processed downstream, the synthetic key-value pairs will trigger evaluation of any
   * pending timers for all keys, even those that haven't received new data in the current batch.
   * This ensures consistent and reliable timer execution across the entire stateful pipeline.
   *
   * <p>This approach is critical for implementing features like windowing, sessions, and other
   * time-based operations that rely on timers firing at the correct moments, regardless of data
   * activity.
   *
   * @param <KeyT> The type of keys
   * @param <ValueT> The type of values
   * @param originalLRDD The original DStream of windowed key-value pairs to be enhanced with
   *     periodic timer events
   * @return A new DStream that includes periodic events for all known keys, ensuring timer
   *     processing for every key in each batch
   */
  @SuppressWarnings({"unchecked", "nullness"})
  public static <KeyT, ValueT> JavaDStream<WindowedValue<KV<KeyT, ValueT>>> toPeriodicDStream(
      JavaDStream<WindowedValue<KV<KeyT, ValueT>>> originalLRDD) {
    // extract only unique keys from original RDD
    final JavaPairDStream<KeyT, KeyT> uniqueKeysDStream =
        originalLRDD.mapPartitionsToPair(
            (Iterator<WindowedValue<KV<KeyT, ValueT>>> iterator) -> {
              final Set<KeyT> uniqueKeys = new HashSet<>();
              while (iterator.hasNext()) {
                final KV<KeyT, ValueT> kv = iterator.next().getValue();
                uniqueKeys.add(kv.getKey());
              }

              return Iterators.transform(uniqueKeys.iterator(), key -> Tuple2.apply(null, key));
            });

    final JavaPairDStream</*(null)*/ KeyT, /*Actual Keys*/ Set<KeyT>> keyStateSnapshotDStream =
        uniqueKeysDStream
            .mapWithState(
                StateSpec.function(
                    new SerializableFunction3<
                        /*(null)*/ KeyT,
                        /*Actual Key*/ Option<KeyT>,
                        /*State*/ State<Set<KeyT>>,
                        /*Return type - not needed as we use stateSnapshots*/ Void>() {
                      @Override
                      public Void apply(
                          KeyT timerKeyMarker, Option<KeyT> actualKey, State<Set<KeyT>> state) {
                        Set<KeyT> keyState;
                        if (state.exists()) {
                          keyState = state.get();
                        } else {
                          keyState = new HashSet<>();
                        }

                        keyState.add(actualKey.get());
                        state.update(keyState);
                        // We don't need to return anything as we use stateSnapshots
                        return null;
                      }
                    }))
            .stateSnapshots();

    return originalLRDD.transformWith(
        keyStateSnapshotDStream,
        (JavaRDD<WindowedValue<KV<KeyT, ValueT>>> wvRDD,
            JavaPairRDD<KeyT, Set<KeyT>> stateRDD,
            Time time) -> {
          final int numPartitions = wvRDD.getNumPartitions();

          final Set<KeyT> keySet = new HashSet<>();

          // add all keys to keySet
          stateRDD.distinct().values().collect().forEach(keySet::addAll);

          final List<WindowedValue<KV<KeyT, ValueT>>> collect =
              keySet.stream()
                  .map(key -> new WindowedValueForTimerMarker<>(KV.of(key, (ValueT) TIMER_MARKER)))
                  .collect(Collectors.toList());

          final JavaRDD<WindowedValue<KV<KeyT, ValueT>>> keyConvertedRDD =
              JavaRDD.fromRDD(
                  stateRDD
                      .context()
                      .parallelize(
                          JavaConverters.asScalaIterator(collect.iterator()).toSeq(),
                          numPartitions > 0 ? numPartitions : 1,
                          JavaSparkContext$.MODULE$.fakeClassTag()),
                  JavaSparkContext$.MODULE$.fakeClassTag());

          return wvRDD.union(keyConvertedRDD);
        });
  }
}
