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

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.ReadyCheckingSideInputReader;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollectionView;

/**
 * A {@link DoFnRunner} that can refuse to process elements that are not ready, instead returning
 * them via the {@link #processElementInReadyWindows(WindowedValue)}.
 */
public class PushbackSideInputDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, OutputT> {
  private final DoFnRunner<InputT, OutputT> underlying;
  private final Collection<PCollectionView<?>> views;
  private final ReadyCheckingSideInputReader sideInputReader;

  private Set<BoundedWindow> notReadyWindows;

  public static <InputT, OutputT> PushbackSideInputDoFnRunner<InputT, OutputT> create(
      DoFnRunner<InputT, OutputT> underlying,
      Collection<PCollectionView<?>> views,
      ReadyCheckingSideInputReader sideInputReader) {
    return new PushbackSideInputDoFnRunner<>(underlying, views, sideInputReader);
  }

  private PushbackSideInputDoFnRunner(
      DoFnRunner<InputT, OutputT> underlying,
      Collection<PCollectionView<?>> views,
      ReadyCheckingSideInputReader sideInputReader) {
    this.underlying = underlying;
    this.views = views;
    this.sideInputReader = sideInputReader;
  }

  @Override
  public void startBundle() {
    notReadyWindows = new HashSet<>();
    underlying.startBundle();
  }

  /**
   * Call the underlying {@link DoFnRunner#processElement(WindowedValue)} for the provided element
   * for each window the element is in that is ready.
   *
   * @param elem the element to process in all ready windows
   * @return each element that could not be processed because it requires a side input window
   * that is not ready.
   */
  public Iterable<WindowedValue<InputT>> processElementInReadyWindows(WindowedValue<InputT> elem) {
    if (views.isEmpty()) {
      processElement(elem);
      return Collections.emptyList();
    }
    ImmutableList.Builder<WindowedValue<InputT>> pushedBack = ImmutableList.builder();
    for (WindowedValue<InputT> windowElem : elem.explodeWindows()) {
      BoundedWindow mainInputWindow = Iterables.getOnlyElement(windowElem.getWindows());
      boolean isReady = !notReadyWindows.contains(mainInputWindow);
      for (PCollectionView<?> view : views) {
        BoundedWindow sideInputWindow =
            view.getWindowingStrategyInternal()
                .getWindowFn()
                .getSideInputWindow(mainInputWindow);
        if (!sideInputReader.isReady(view, sideInputWindow)) {
          isReady = false;
          break;
        }
      }
      if (isReady) {
        processElement(windowElem);
      } else {
        notReadyWindows.add(mainInputWindow);
        pushedBack.add(windowElem);
      }
    }
    return pushedBack.build();
  }

  @Override
  public void processElement(WindowedValue<InputT> elem) {
    underlying.processElement(elem);
  }

  /**
   * Call the underlying {@link DoFnRunner#finishBundle()}.
   */
  @Override
  public void finishBundle() {
    notReadyWindows = null;
    underlying.finishBundle();
  }
}


@@ -1037,17 +1037,16 @@
 il.Set;%0A
-%0A
 import o
@@ -3283,27 +3283,290 @@
     
-boolean isReady = !
+if (isReady(mainInputWindow)) %7B%0A        processElement(windowElem);%0A      %7D else %7B%0A        notReadyWindows.add(mainInputWindow);%0A        pushedBack.add(windowElem);%0A      %7D%0A    %7D%0A    return pushedBack.build();%0A  %7D%0A%0A  private boolean isReady(BoundedWindow mainInputWindow) %7B%0A    if (
 notR
@@ -3598,28 +3598,54 @@
 InputWindow)
-;%0A  
+) %7B%0A      return false;%0A    %7D%0A
     for (PCo
@@ -3678,26 +3678,24 @@
 ws) %7B%0A      
- 
- 
 BoundedWindo
@@ -3714,18 +3714,16 @@
 indow =%0A
-  
         
@@ -3763,56 +3763,22 @@
 al()
-%0A                .getWindowFn()%0A                
+.getWindowFn()
 .get
@@ -3807,26 +3807,24 @@
 putWindow);%0A
-  
       if (!s
@@ -3884,214 +3884,20 @@
     
-  isReady = false;%0A          break;%0A        %7D%0A      %7D%0A      if (isReady) %7B%0A        processElement(windowElem);%0A      %7D else %7B%0A        notReadyWindows.add(mainInputWindow);%0A        pushedBack.add(windowElem)
+return false
 ;%0A  
@@ -3919,34 +3919,20 @@
  return 
-pushedBack.build()
+true
 ;%0A  %7D%0A%0A 

Remove Window.Unbound