/*
 * 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.flink.table.runtime.operators.window.groupwindow.internal;

import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.runtime.generated.NamespaceAggsHandleFunctionBase;
import org.apache.flink.table.runtime.operators.window.Window;
import org.apache.flink.table.runtime.operators.window.groupwindow.assigners.MergingWindowAssigner;
import org.apache.flink.util.function.BiConsumerWithException;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import static org.apache.flink.table.runtime.util.TimeWindowUtil.toEpochMillsForTimer;
import static org.apache.flink.util.Preconditions.checkArgument;

/**
 * The implementation of {@link InternalWindowProcessFunction} for {@link MergingWindowAssigner}.
 *
 * @param <W> The type of {@code Window} that assigner assigns.
 */
public class MergingWindowProcessFunction<K, W extends Window>
        extends InternalWindowProcessFunction<K, W> {

    private static final long serialVersionUID = -2866771637946397223L;

    private final MergingWindowAssigner<W> windowAssigner;
    private final TypeSerializer<W> windowSerializer;
    private transient MergingWindowSet<W> mergingWindows;
    private transient MergingFunctionImpl mergingFunction;

    private List<W> reuseActualWindows;

    public MergingWindowProcessFunction(
            MergingWindowAssigner<W> windowAssigner,
            NamespaceAggsHandleFunctionBase<W> windowAggregator,
            TypeSerializer<W> windowSerializer,
            long allowedLateness) {
        super(windowAssigner, windowAggregator, allowedLateness);
        this.windowAssigner = windowAssigner;
        this.windowSerializer = windowSerializer;
    }

    @Override
    public void open(Context<K, W> ctx) throws Exception {
        checkArgument(ctx instanceof MergingContext);
        super.open(ctx);
        MapStateDescriptor<W, W> mappingStateDescriptor =
                new MapStateDescriptor<>(
                        "session-window-mapping", windowSerializer, windowSerializer);
        MapState<W, W> windowMapping = ctx.getPartitionedState(mappingStateDescriptor);
        this.mergingWindows = new MergingWindowSet<>(windowAssigner, windowMapping);
        this.mergingFunction =
                new MergingFunctionImpl(((MergingContext) ctx).getWindowStateMergingConsumer());
    }

    @Override
    public Collection<W> assignStateNamespace(RowData inputRow, long timestamp) throws Exception {
        Collection<W> elementWindows = windowAssigner.assignWindows(inputRow, timestamp);
        mergingWindows.initializeCache(ctx.currentKey());
        reuseActualWindows = new ArrayList<>(1);
        for (W window : elementWindows) {
            // adding the new window might result in a merge, in that case the actualWindow
            // is the merged window and we work with that. If we don't merge then
            // actualWindow == window
            W actualWindow = mergingWindows.addWindow(window, mergingFunction);

            // drop if the window is already late
            if (isWindowLate(actualWindow)) {
                mergingWindows.retireWindow(actualWindow);
            } else {
                reuseActualWindows.add(actualWindow);
            }
        }
        List<W> affectedWindows = new ArrayList<>(reuseActualWindows.size());
        for (W actual : reuseActualWindows) {
            affectedWindows.add(mergingWindows.getStateWindow(actual));
        }
        return affectedWindows;
    }

    @Override
    public Collection<W> assignActualWindows(RowData inputRow, long timestamp) throws Exception {
        // the actual windows is calculated in assignStateNamespace
        return reuseActualWindows;
    }

    @Override
    public void prepareAggregateAccumulatorForEmit(W window) throws Exception {
        W stateWindow = mergingWindows.getStateWindow(window);
        RowData acc = ctx.getWindowAccumulators(stateWindow);
        if (acc == null) {
            acc = windowAggregator.createAccumulators();
        }
        windowAggregator.setAccumulators(stateWindow, acc);
    }

    @Override
    public void cleanWindowIfNeeded(W window, long currentTime) throws Exception {
        if (isCleanupTime(window, currentTime)) {
            ctx.clearTrigger(window);
            W stateWindow = mergingWindows.getStateWindow(window);
            ctx.clearWindowState(stateWindow);
            // retire expired window
            mergingWindows.initializeCache(ctx.currentKey());
            mergingWindows.retireWindow(window);
            // do not need to clear previous state, previous state is disabled in session window
        }
    }

    /** Get the state window as the namespace stored acc in the data about this actual window. */
    public W getStateWindow(W window) throws Exception {
        return mergingWindows.getStateWindow(window);
    }

    private class MergingFunctionImpl implements MergingWindowSet.MergeFunction<W> {

        private final BiConsumerWithException<W, Collection<W>, Throwable> accMergingConsumer;

        public MergingFunctionImpl(
                BiConsumerWithException<W, Collection<W>, Throwable> accMergingConsumer) {
            this.accMergingConsumer = accMergingConsumer;
        }

        @Override
        public void merge(
                W mergeResult,
                Collection<W> mergedWindows,
                W stateWindowResult,
                Collection<W> stateWindowsToBeMerged)
                throws Exception {

            long mergeResultMaxTs =
                    toEpochMillsForTimer(mergeResult.maxTimestamp(), ctx.getShiftTimeZone());
            if ((windowAssigner.isEventTime()
                    && mergeResultMaxTs + allowedLateness <= ctx.currentWatermark())) {
                throw new UnsupportedOperationException(
                        "The end timestamp of an "
                                + "event-time window cannot become earlier than the current watermark "
                                + "by merging. Current watermark: "
                                + ctx.currentWatermark()
                                + " window: "
                                + mergeResult);
            } else if (!windowAssigner.isEventTime()
                    && mergeResultMaxTs <= ctx.currentProcessingTime()) {
                throw new UnsupportedOperationException(
                        "The end timestamp of a "
                                + "processing-time window cannot become earlier than the current processing time "
                                + "by merging. Current processing time: "
                                + ctx.currentProcessingTime()
                                + " window: "
                                + mergeResult);
            }

            ctx.onMerge(mergeResult, stateWindowsToBeMerged);

            // clear registered timers
            for (W m : mergedWindows) {
                ctx.clearTrigger(m);
                ctx.deleteCleanupTimer(m);
            }

            // merge the merged state windows into the newly resulting state window
            if (!stateWindowsToBeMerged.isEmpty()) {
                try {
                    accMergingConsumer.accept(stateWindowResult, stateWindowsToBeMerged);
                } catch (Throwable e) {
                    throw new RuntimeException("Should not happen", e);
                }
            }
        }
    }

    /**
     * A default merging consumer that merges the accumulators in state windows that waiting to be
     * merged into the target state window.
     *
     * <p>The first parameter is the target state window, the second parameter is the windows that
     * need to be merged.
     */
    public static class DefaultAccMergingConsumer<W extends Window>
            implements BiConsumerWithException<W, Collection<W>, Throwable> {

        private final Context<?, W> ctx;

        private final NamespaceAggsHandleFunctionBase<W> windowAggregator;

        public DefaultAccMergingConsumer(
                Context<?, W> ctx, NamespaceAggsHandleFunctionBase<W> windowAggregator) {
            this.ctx = ctx;
            this.windowAggregator = windowAggregator;
        }

        @Override
        public void accept(W stateWindowResult, Collection<W> stateWindowsToBeMerged)
                throws Throwable {
            RowData targetAcc = ctx.getWindowAccumulators(stateWindowResult);
            if (targetAcc == null) {
                targetAcc = windowAggregator.createAccumulators();
            }
            windowAggregator.setAccumulators(stateWindowResult, targetAcc);
            for (W w : stateWindowsToBeMerged) {
                RowData acc = ctx.getWindowAccumulators(w);
                if (acc != null) {
                    windowAggregator.merge(w, acc);
                }
                // clear merged window
                ctx.clearWindowState(w);
                ctx.clearPreviousState(w);
            }
            targetAcc = windowAggregator.getAccumulators();
            ctx.setWindowAccumulators(stateWindowResult, targetAcc);
        }
    }

    /** A {@link Context} used for {@link MergingWindowProcessFunction}. */
    public interface MergingContext<K, W extends Window> extends Context<K, W> {

        /** Returns the consumer used to merge window state. */
        BiConsumerWithException<W, Collection<W>, Throwable> getWindowStateMergingConsumer();
    }
}
