public interface TableWriteFunction<K, V, U> extends TableFunction {
  /**
   * Asynchronously update the record with specified {@code key} and additional arguments.
   * This method must be thread-safe.
   *
   * If the update operation failed due to the an existing record missing for the key, the implementation can return
   * a future completed exceptionally with a {@link RecordNotFoundException} which will
   * allow to Put a default value if one is provided.
   *
   * @param key key for the table record
   * @param update update record for the given key
   * @return CompletableFuture for the update request
   */
  CompletableFuture<Void> updateAsync(K key, U update);

  /**
   * Asynchronously update the record with specified {@code key} and additional arguments.
   * This method must be thread-safe.
   *
   * If the update operation failed due to the an existing record missing for the key, the implementation can return
   * a future completed exceptionally with a {@link RecordNotFoundException} which will
   * allow to Put a default value if one is provided.
   *
   * @param key key for the table record
   * @param update update record for the given key
   * @param args additional arguments
   * @return CompletableFuture for the update request
   */
  default CompletableFuture<Void> updateAsync(K key, U update, Object ... args) {
    throw new SamzaException("Not supported");
  }

  /**
   * Asynchronously updates the table with {@code records} with specified {@code keys}. This method must be thread-safe.
   * The default implementation calls updateAsync for each entry and return a combined future.
   * @param records updates for the table
   * @return CompletableFuture for the update request
   */
  default CompletableFuture<Void> updateAllAsync(Collection<Entry<K, U>> records) {
    final List<CompletableFuture<Void>> updateFutures = records.stream()
        .map(entry -> updateAsync(entry.getKey(), entry.getValue()))
        .collect(Collectors.toList());
    return CompletableFuture.allOf(Iterables.toArray(updateFutures, CompletableFuture.class));
  }
}

public interface MessageStream<M> {
/**
  * Allows sending update messages in this {@link MessageStream} to a {@link Table} and then propagates this
  * {@link MessageStream} to the next chained operator. The type of input message is expected to be {@link KV},
  * otherwise a {@link ClassCastException} will be thrown. The value is an UpdateMessage- update and default value.
  * Defaults are optional and can be used if the Remote Table integration supports inserting a default through PUT in
  * the event an update fails due to an existing record being absent.
  * <p>
  * Note: The update will be written but may not be flushed to the underlying table before its propagated to the
  * chained operators. Whether the message can be read back from the Table in the chained operator depends on whether
  * it was flushed and whether the Table offers read after write consistency. Messages retain the original partitioning
  * scheme when propagated to next operator.
  *
  * @param table the table to write messages to
  * @param contract Update contract which defines how the update will be performed   
  * @param <K> the type of key in the table
  * @param <V> the type of record value in the table
  * @param <U> the type of update value for the table
  * @return this {@link MessageStream}
  */
 <K, V, U> MessageStream<KV<K, UpdateMessage<U, V>>> sendTo(Table<KV<K, V>> table, UpdateContract contract);
}

/**
 * Represents an update and an optional default record to be inserted for a key,
 * if the update is applied to a non-existent record.
 *
 * @param <U> type of the update record
 * @param <V> type of the default record
 */
public final class UpdateMessage<U, V> {
  private final U update;
  @Nullable private final V defaultValue;

  public static <U, V> UpdateMessage<U, V> of(U update, @Nullable V defaultValue) {
    return new UpdateMessage<>(update, defaultValue);
  }

  public static <U, V> UpdateMessage<U, V> of(U update) {
    return new UpdateMessage<>(update, null);
  }

  private UpdateMessage(U update, V defaultValue) {
    this.update = update;
    this.defaultValue = defaultValue;
  }

  public U getUpdate() {
    return update;
  }

  public V getDefault() {
    return defaultValue;
  }
}

public class SendToTableWithUpdateOperatorImpl<K, V, U>
    extends OperatorImpl<KV<K, UpdateMessage<U, V>>, KV<K, UpdateMessage<U, V>>> {
  private static final Logger LOG = LoggerFactory.getLogger(SendToTableWithUpdateOperatorImpl.class);

  private final SendToTableWithUpdateOperatorSpec<K, V, U> spec;
  private final ReadWriteTable<K, V, U> table;

  public SendToTableWithUpdateOperatorImpl(SendToTableWithUpdateOperatorSpec<K, V, U> spec, Context context) {
    this.spec = spec;
    this.table = context.getTaskContext().getTable(spec.getTableId());
    if (context.getTaskContext().getTable(spec.getTableId()) instanceof RemoteTable) {
      RemoteTable<K, V, U> remoteTable = (RemoteTable<K, V, U>) table;
      if (remoteTable.getBatchProvider() instanceof CompactBatchProvider) {
        throw new SamzaException("Batching is not supported with Compact Batches for partial updates");
      }
    }
  }

  @Override
  protected void handleInit(Context context) {
  }

  @Override
  protected CompletionStage<Collection<KV<K, UpdateMessage<U, V>>>> handleMessageAsync(KV<K, UpdateMessage<U, V>> message,
      MessageCollector collector, TaskCoordinator coordinator) {
    final UpdateContract contract = spec.getContract();
    final CompletableFuture<Void> updateFuture = table.updateAsync(message.getKey(), message.getValue().getUpdate());

    return updateFuture
        .handle((result, ex) -> {
          if (ex == null) {
            // success, no need to Put a default value
            return false;
          } else if (ex.getCause() instanceof RecordNotFoundException && message.getValue().hasDefault()) {
            // If an update fails for a given key due to a RecordDoesNotExistException exception thrown and
            // a default is provided and the UpdateContract is set to UPDATE_WITH_DEFAULTS, then attempt
            // to PUT a default record for the key and then apply the update.
            if (contract == UpdateContract.UPDATE_WITH_DEFAULTS) {
              return true;
            } else {
              throw new SamzaException("Put default failed for update as the Update contract was set to " + contract +
                  ". Please use UpdateContract.UPDATE_WITH_DEFAULTS instead.");
            }
          } else {
            throw new SamzaException("Update failed with exception: ", ex);
          }
        })
        .thenCompose(shouldPutDefault -> {
          if (shouldPutDefault) {
            final CompletableFuture<Void> putFuture = table.putAsync(message.getKey(), message.getValue().getDefault());
            return putFuture
                .exceptionally(ex -> {
                  LOG.warn("PUT default failed due to an exception. Ignoring the exception and proceeding with update. "
                          + "The exception encountered is: ", ex);
                  return null;
                })
                .thenCompose(res -> table.updateAsync(message.getKey(), message.getValue().getUpdate()));
          } else {
            return CompletableFuture.completedFuture(null);
          }
        }).thenApply(result -> Collections.singleton(message));
  }

  @Override
  protected void handleClose() {
    table.close();
  }

  @Override
  protected OperatorSpec<KV<K, UpdateMessage<U, V>>, KV<K, UpdateMessage<U, V>>> getOperatorSpec() {
    return spec;
  }
}  

/**
 * Interface for table operations that can be batched.
 *
 * @param <K> The key type associated with the operation.
 * @param <V> The value type associated with the operation.
 * @param <U> The update type associated with the operation.
 */
public interface Operation<K, V, U> {
  /**
   * @return The key associated with the operation.
   */
  K getKey();

  /**
   * @return The value associated with the operation.
   */
  V getValue();

  /**
   * @return The update associated with the operation.
   */
  U getUpdate();

  /**
   * @return The extra arguments associated with the operation.
   */
  Object[] getArgs();
}

/**
 * Update operation.
 *
 * @param <K> The type of the key.
 * @param <U> The type of the update
 */
public class UpdateOperation<K, V, U> implements Operation<K, V, U> {
  final private K key;
  final private U update;
  final private Object[] args;

  public UpdateOperation(K key, U update, Object ... args) {
    Preconditions.checkNotNull(key);
    Preconditions.checkNotNull(update);
    Preconditions.checkNotNull(args);
    this.key = key;
    this.update = update;
    this.args = args;
  }

  /**
   * @return The key to be updated in the table.
   */
  @Override
  public K getKey() {
    return key;
  }

  /**
   * @return null.
   */
  @Override
  public V getValue() {
    return null;
  }

  /**
   * @return The Update to be applied to the table for the key.
   */
  @Override
  public U getUpdate() {
    return update;
  }

  @Override
  public Object[] getArgs() {
    return args;
  }
}