/*
 * Copyright 2017 Spotify AB.
 *
 * Licensed 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 com.spotify.scio.extra.transforms;

import com.google.common.collect.Maps;
import com.google.common.collect.Queues;
import org.apache.beam.sdk.transforms.DoFn;

import java.util.UUID;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.function.Function;

/**
 * A {@link DoFn} that handles asynchronous requests to an external service.
 */
public abstract class BaseAsyncDoFn<InputT, OutputT, ResourceT, FutureT>
    extends DoFnWithResource<InputT, OutputT, ResourceT> {

  /**
   * Process an element asynchronously.
   */
  public abstract FutureT processElement(InputT input);

  protected abstract void waitForFutures(Iterable<FutureT> futures)
      throws InterruptedException, ExecutionException;
  protected abstract FutureT addCallback(FutureT future,
                                         Function<OutputT, Void> onSuccess,
                                         Function<Throwable, Void> onFailure);

  private final ConcurrentMap<UUID, FutureT> futures = Maps.newConcurrentMap();
  private final ConcurrentLinkedQueue<OutputT> results = Queues.newConcurrentLinkedQueue();
  private final ConcurrentLinkedQueue<Throwable> errors = Queues.newConcurrentLinkedQueue();

  @StartBundle
  public void startBundle(StartBundleContext c) {
    futures.clear();
    results.clear();
    errors.clear();
  }

  @FinishBundle
  public void finishBundle(ProcessContext c) {
    if (!futures.isEmpty()) {
      try {
        waitForFutures(futures.values());
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        throw new RuntimeException("Failed to process futures", e);
      } catch (ExecutionException e) {
        throw new RuntimeException("Failed to process futures", e);
      }
    }
    flush(c);
  }

  @ProcessElement
  public void processElement(ProcessContext c) {
    flush(c);

    final UUID uuid = UUID.randomUUID();
    FutureT future = addCallback(processElement(c.element()), r -> {
      results.add(r);
      futures.remove(uuid);
      return null;
    }, t -> {
      errors.add(t);
      futures.remove(uuid);
      return null;
    });
    // This `put` may happen after `remove` in the callbacks but it's OK since either the result
    // or the error would've already been pushed to the corresponding queues and we are not losing
    // data. `waitForFutures` in `finishBundle` blocks until all pending futures, including ones
    // that may have already completed, and `startBundle` clears everything.
    futures.put(uuid, future);
  }

  private void flush(ProcessContext c) {
    if (!errors.isEmpty()) {
      RuntimeException e = new RuntimeException("Failed to process futures");
      Throwable t = errors.poll();
      while (t != null) {
        e.addSuppressed(t);
        t = errors.poll();
      }
      throw e;
    }
    OutputT element = results.poll();
    while (element != null) {
      c.output(element);
      element = results.poll();
    }
  }

}

@@ -760,16 +760,109 @@
 ms.DoFn;
+%0Aimport org.apache.beam.sdk.transforms.windowing.BoundedWindow;%0Aimport org.joda.time.Instant;
 %0A%0Aimport
@@ -1838,23 +1838,22 @@
 edQueue%3C
-OutputT
+Result
 %3E result
@@ -2030,28 +2030,8 @@
 dle(
-StartBundleContext c
 ) %7B%0A
@@ -2140,23 +2140,28 @@
 hBundle(
-Process
+FinishBundle
 Context 
@@ -2597,32 +2597,54 @@
 ProcessContext c
+, BoundedWindow window
 ) %7B%0A    flush(c)
@@ -2774,17 +2774,52 @@
 lts.add(
-r
+new Result(r, c.timestamp(), window)
 );%0A     
@@ -3657,32 +3657,300 @@
     
-OutputT element = result
+Result r = results.poll();%0A    while (r != null) %7B%0A      c.output(r.output);%0A      r = results.poll();%0A    %7D%0A  %7D%0A%0A  private void flush(FinishBundleContext c) %7B%0A    if (!errors.isEmpty()) %7B%0A      RuntimeException e = new RuntimeException(%22Failed to process futures%22);%0A      Throwable t = error
 s.po
@@ -3963,89 +3963,532 @@
     
+  
 while (
-element != null) %7B%0A      c.output(element);%0A      element = results.poll()
+t != null) %7B%0A        e.addSuppressed(t);%0A        t = errors.poll();%0A      %7D%0A      throw e;%0A    %7D%0A    Result r = results.poll();%0A    while (r != null) %7B%0A      c.output(r.output, r.timestamp, r.window);%0A      r = results.poll();%0A    %7D%0A  %7D%0A%0A  private class Result %7B%0A    private OutputT output;%0A    private Instant timestamp;%0A    private BoundedWindow window;%0A%0A    Result(OutputT output, Instant timestamp, BoundedWindow window) %7B%0A      this.output = output;%0A      this.timestamp = timestamp;%0A      this.window = window
 ;%0A  

fix BaseAsyncDoFn for Beam 2.0