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

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.beam.runners.core.DoFnRunner;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.StateNamespaces.WindowNamespace;
import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.TimerInternals.TimerData;
import org.apache.beam.runners.core.TimerInternals.TimerDataCoder;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.MapCoder;
import org.apache.beam.sdk.coders.SetCoder;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.state.WatermarkHoldState;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.Parser;
import org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Lists;

/** A class that handles streaming side inputs in a {@link DoFnRunner}. */
public class StreamingSideInputFetcher<InputT, W extends BoundedWindow> {
  private StreamingModeExecutionContext.StreamingModeStepContext stepContext;
  private Map<String, PCollectionView<?>> sideInputViews;

  private final StateTag<BagState<WindowedValue<InputT>>> elementsAddr;
  private final StateTag<BagState<TimerData>> timersAddr;
  private final StateTag<WatermarkHoldState> watermarkHoldingAddr;
  private final StateTag<ValueState<Map<W, Set<Windmill.GlobalDataRequest>>>> blockedMapAddr;

  private Map<W, Set<Windmill.GlobalDataRequest>> blockedMap = null; // lazily initialized

  private final Coder<W> mainWindowCoder;

  public StreamingSideInputFetcher(
      Iterable<PCollectionView<?>> views,
      Coder<InputT> inputCoder,
      WindowingStrategy<?, W> windowingStrategy,
      StreamingModeExecutionContext.StreamingModeStepContext stepContext) {
    this.stepContext = stepContext;

    this.mainWindowCoder = windowingStrategy.getWindowFn().windowCoder();

    this.sideInputViews = new HashMap<>();
    for (PCollectionView<?> view : views) {
      sideInputViews.put(view.getTagInternal().getId(), view);
    }

    this.blockedMapAddr = blockedMapAddr(mainWindowCoder);
    this.elementsAddr =
        StateTags.makeSystemTagInternal(
            StateTags.bag("elem", WindowedValue.getFullCoder(inputCoder, mainWindowCoder)));
    this.timersAddr =
        StateTags.makeSystemTagInternal(StateTags.bag("timer", TimerDataCoder.of(mainWindowCoder)));
    StateTag<WatermarkHoldState> watermarkTag =
        StateTags.watermarkStateInternal(
            "holdForSideinput", windowingStrategy.getTimestampCombiner());
    this.watermarkHoldingAddr = StateTags.makeSystemTagInternal(watermarkTag);
  }

  @VisibleForTesting
  static <W extends BoundedWindow>
      StateTag<ValueState<Map<W, Set<GlobalDataRequest>>>> blockedMapAddr(
          Coder<W> mainWindowCoder) {
    return StateTags.value(
        "blockedMap", MapCoder.of(mainWindowCoder, SetCoder.of(GlobalDataRequestCoder.of())));
  }

  /** Computes the set of main input windows for which all side inputs are ready and cached. */
  public Set<W> getReadyWindows() {
    Set<W> readyWindows = new HashSet<>();

    for (Windmill.GlobalDataId id : stepContext.getSideInputNotifications()) {
      if (sideInputViews.get(id.getTag()) == null) {
        // Side input is for a different DoFn; ignore it.
        continue;
      }

      for (Map.Entry<W, Set<Windmill.GlobalDataRequest>> entry : blockedMap().entrySet()) {
        Set<Windmill.GlobalDataRequest> windowBlockedSet = entry.getValue();
        Set<Windmill.GlobalDataRequest> found = new HashSet<>();
        for (Windmill.GlobalDataRequest request : windowBlockedSet) {
          if (id.equals(request.getDataId())) {
            found.add(request);
          }
        }

        windowBlockedSet.removeAll(found);

        if (windowBlockedSet.isEmpty()) {
          // Notifications were received for all side inputs for this window.
          // Issue fetches for all the needed side inputs to make sure they are all present
          // in the local cache.  If not, note the side inputs as still being blocked.
          W window = entry.getKey();
          boolean allSideInputsCached = true;
          for (PCollectionView<?> view : sideInputViews.values()) {
            if (!stepContext.issueSideInputFetch(
                view, window, StateFetcher.SideInputState.KNOWN_READY)) {
              Windmill.GlobalDataRequest request = buildGlobalDataRequest(view, window);
              stepContext.addBlockingSideInput(request);
              windowBlockedSet.add(request);
              allSideInputsCached = false;
            }
          }

          if (allSideInputsCached) {
            readyWindows.add(window);
          }
        }
      }
    }

    return readyWindows;
  }

  public Iterable<BagState<WindowedValue<InputT>>> prefetchElements(Iterable<W> readyWindows) {
    List<BagState<WindowedValue<InputT>>> elements = Lists.newArrayList();
    for (W window : readyWindows) {
      elements.add(elementBag(window).readLater());
    }
    return elements;
  }

  public void releaseBlockedWindows(Iterable<W> windows) {
    for (W window : windows) {
      WatermarkHoldState watermarkHold = watermarkHold(window);
      watermarkHold.clear();

      blockedMap().remove(window);
    }
  }

  public Iterable<BagState<TimerData>> prefetchTimers(Iterable<W> readyWindows) {
    List<BagState<TimerData>> timers = Lists.newArrayList();
    for (W window : readyWindows) {
      timers.add(timerBag(window).readLater());
    }
    return timers;
  }

  /** Compute the set of side inputs that are not yet ready for the given main input window. */
  public boolean storeIfBlocked(WindowedValue<InputT> elem) {
    @SuppressWarnings("unchecked")
    W window = (W) Iterables.getOnlyElement(elem.getWindows());

    Set<Windmill.GlobalDataRequest> blocked = blockedMap().get(window);
    if (blocked == null) {
      for (PCollectionView<?> view : sideInputViews.values()) {
        if (!stepContext.issueSideInputFetch(view, window, StateFetcher.SideInputState.UNKNOWN)) {
          if (blocked == null) {
            blocked = new HashSet<>();
            blockedMap().put(window, blocked);
          }
          blocked.add(buildGlobalDataRequest(view, window));
        }
      }
    }
    if (blocked != null) {
      elementBag(window).add(elem);
      watermarkHold(window).add(elem.getTimestamp());
      stepContext.addBlockingSideInputs(blocked);
      return true;
    } else {
      return false;
    }
  }

  public boolean storeIfBlocked(TimerData timer) {
    if (!(timer.getNamespace() instanceof WindowNamespace)) {
      throw new IllegalArgumentException(
          "Expected WindowNamespace, but was " + timer.getNamespace());
    }
    @SuppressWarnings("unchecked")
    WindowNamespace<W> windowNamespace = (WindowNamespace<W>) timer.getNamespace();
    W window = windowNamespace.getWindow();

    boolean blocked = false;
    for (PCollectionView<?> view : sideInputViews.values()) {
      if (!stepContext.issueSideInputFetch(view, window, StateFetcher.SideInputState.UNKNOWN)) {
        blocked = true;
      }
    }
    if (blocked) {
      timerBag(window).add(timer);
    }
    return blocked;
  }

  public void persist() {
    if (blockedMap == null) {
      return;
    }

    ValueState<Map<W, Set<Windmill.GlobalDataRequest>>> mapState =
        stepContext.stateInternals().state(StateNamespaces.global(), blockedMapAddr);
    if (blockedMap.isEmpty()) {
      // Avoid storing the empty map so we don't leave unnecessary state behind from processing
      // the key.
      mapState.clear();
    } else {
      mapState.write(blockedMap);
    }
    blockedMap = null;
  }

  private Map<W, Set<Windmill.GlobalDataRequest>> blockedMap() {
    if (blockedMap == null) {
      blockedMap =
          stepContext.stateInternals().state(StateNamespaces.global(), blockedMapAddr).read();
      if (blockedMap == null) {
        blockedMap = new HashMap<>();
      }
    }
    return blockedMap;
  }

  @VisibleForTesting
  Set<W> getBlockedWindows() {
    return blockedMap().keySet();
  }

  @VisibleForTesting
  BagState<WindowedValue<InputT>> elementBag(W window) {
    return stepContext
        .stateInternals()
        .state(StateNamespaces.window(mainWindowCoder, window), elementsAddr);
  }

  @VisibleForTesting
  WatermarkHoldState watermarkHold(W window) {
    return stepContext
        .stateInternals()
        .state(StateNamespaces.window(mainWindowCoder, window), watermarkHoldingAddr);
  }

  @VisibleForTesting
  BagState<TimerData> timerBag(W window) {
    return stepContext
        .stateInternals()
        .state(StateNamespaces.window(mainWindowCoder, window), timersAddr);
  }

  private <SideWindowT extends BoundedWindow> Windmill.GlobalDataRequest buildGlobalDataRequest(
      PCollectionView<?> view, BoundedWindow mainWindow) {
    @SuppressWarnings("unchecked")
    WindowingStrategy<?, SideWindowT> sideWindowStrategy =
        (WindowingStrategy<?, SideWindowT>) view.getWindowingStrategyInternal();

    WindowFn<?, SideWindowT> sideWindowFn = sideWindowStrategy.getWindowFn();

    Coder<SideWindowT> sideInputWindowCoder = sideWindowFn.windowCoder();

    SideWindowT sideInputWindow =
        (SideWindowT) view.getWindowMappingFn().getSideInputWindow(mainWindow);

    ByteString.Output windowStream = ByteString.newOutput();
    try {
      sideInputWindowCoder.encode(sideInputWindow, windowStream, Coder.Context.OUTER);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

    return Windmill.GlobalDataRequest.newBuilder()
        .setDataId(
            Windmill.GlobalDataId.newBuilder()
                .setTag(view.getTagInternal().getId())
                .setVersion(windowStream.toByteString())
                .build())
        .setExistenceWatermarkDeadline(
            WindmillTimeUtils.harnessToWindmillTimestamp(
                sideWindowStrategy.getTrigger().getWatermarkThatGuaranteesFiring(sideInputWindow)))
        .build();
  }

  private static class GlobalDataRequestCoder extends AtomicCoder<GlobalDataRequest> {
    private final Class<Windmill.GlobalDataRequest> protoMessageClass =
        Windmill.GlobalDataRequest.class;
    private transient Parser<Windmill.GlobalDataRequest> memoizedParser;

    public static GlobalDataRequestCoder of() {
      return new GlobalDataRequestCoder();
    }

    @Override
    public Windmill.GlobalDataRequest decode(InputStream inStream) throws IOException {
      return decode(inStream, Context.NESTED);
    }

    @Override
    public Windmill.GlobalDataRequest decode(InputStream inStream, Context context)
        throws IOException {
      if (context.isWholeStream) {
        return getParser().parseFrom(inStream);
      } else {
        return getParser().parseDelimitedFrom(inStream);
      }
    }

    @Override
    public void encode(Windmill.GlobalDataRequest value, OutputStream outStream)
        throws IOException {
      encode(value, outStream, Context.NESTED);
    }

    @Override
    public void encode(Windmill.GlobalDataRequest value, OutputStream outStream, Context context)
        throws IOException {
      if (value == null) {
        throw new CoderException("cannot encode a null " + protoMessageClass.getSimpleName());
      }
      if (context.isWholeStream) {
        value.writeTo(outStream);
      } else {
        value.writeDelimitedTo(outStream);
      }
    }

    private Parser<Windmill.GlobalDataRequest> getParser() {
      if (memoizedParser == null) {
        memoizedParser = Windmill.GlobalDataRequest.getDefaultInstance().getParserForType();
      }
      return memoizedParser;
    }
  }
}