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

import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.OpenContext;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.functions.util.FunctionUtils;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.runtime.generated.FilterCondition;
import org.apache.flink.table.runtime.generated.GeneratedFilterCondition;
import org.apache.flink.table.runtime.generated.GeneratedFunction;
import org.apache.flink.table.runtime.operators.join.FlinkJoinType;
import org.apache.flink.table.runtime.typeutils.RowDataSerializer;
import org.apache.flink.util.Preconditions;

import javax.annotation.Nullable;

import java.io.Serializable;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

/**
 * A runtime delta join tree converted from {@code DeltaJoinTree}.
 *
 * <p>Its structure is fundamentally similar to {@code DeltaJoinTree}. The only difference is that
 * this tree resides within the runtime module, and holds the function generated by the calc on the
 * {@link Node}.
 *
 * <p>See {@code DeltaJoinTree} for more details.
 */
public class DeltaJoinRuntimeTree implements Serializable {

    public final Node root;

    public DeltaJoinRuntimeTree(Node root) {
        this.root = root;

        // validate
        List<Integer> allInputOrdinals =
                root.getAllInputOrdinals().stream().sorted().collect(Collectors.toList());
        Preconditions.checkArgument(
                allInputOrdinals.equals(
                        IntStream.range(0, allInputOrdinals.size())
                                .boxed()
                                .collect(Collectors.toList())));
    }

    public void open(RuntimeContext runtimeContext, OpenContext openContext) throws Exception {
        root.open(runtimeContext, openContext);
    }

    public void close() throws Exception {
        root.close();
    }

    public DeltaJoinRuntimeTree copy() {
        Node newRoot = root.copy();
        return new DeltaJoinRuntimeTree(newRoot);
    }

    /**
     * An abstract node for {@link BinaryInputNode} and {@link JoinNode}.
     *
     * <p>See {@code DeltaJoinTree#Node} for more details.
     */
    public abstract static class Node implements Serializable {

        protected final @Nullable GeneratedFunction<FlatMapFunction<RowData, RowData>>
                generatedCalc;
        public final RowDataSerializer rowDataSerializerPassThroughCalc;

        public transient @Nullable FlatMapFunction<RowData, RowData> calc;

        public Node(
                @Nullable GeneratedFunction<FlatMapFunction<RowData, RowData>> generatedCalc,
                RowDataSerializer rowDataSerializerPassThroughCalc) {
            this.generatedCalc = generatedCalc;
            this.rowDataSerializerPassThroughCalc = rowDataSerializerPassThroughCalc;
        }

        public void open(RuntimeContext runtimeContext, OpenContext openContext) throws Exception {
            if (generatedCalc != null) {
                calc = generatedCalc.newInstance(runtimeContext.getUserCodeClassLoader());
                FunctionUtils.setFunctionRuntimeContext(calc, runtimeContext);
                FunctionUtils.openFunction(calc, openContext);
            }
        }

        public void close() throws Exception {
            if (calc != null) {
                FunctionUtils.closeFunction(calc);
            }
        }

        public abstract Set<Integer> getAllInputOrdinals();

        public abstract Node copy();
    }

    /**
     * A {@link Node} in the tree representing a source used for delta join to scan and lookup.
     *
     * <p>See {@code DeltaJoinTree#BinaryInputNode} for more details.
     */
    public static class BinaryInputNode extends Node {
        public final int inputOrdinal;

        public BinaryInputNode(
                int inputOrdinal,
                @Nullable GeneratedFunction<FlatMapFunction<RowData, RowData>> generatedCalc,
                RowDataSerializer rowDataSerializerPassThroughCalc) {
            super(generatedCalc, rowDataSerializerPassThroughCalc);
            this.inputOrdinal = inputOrdinal;
        }

        @Override
        public Set<Integer> getAllInputOrdinals() {
            return Collections.singleton(inputOrdinal);
        }

        @Override
        public Node copy() {
            return new BinaryInputNode(
                    inputOrdinal, generatedCalc, rowDataSerializerPassThroughCalc);
        }
    }

    /**
     * A {@link Node} in the tree representing a join operation between two inputs.
     *
     * <p>See {@code DeltaJoinTree#JoinNode} for more details.
     */
    public static class JoinNode extends Node {
        public final FlinkJoinType joinType;
        public final GeneratedFilterCondition generatedJoinCondition;
        public final Node left;
        public final Node right;

        public transient FilterCondition joinCondition;

        public JoinNode(
                FlinkJoinType joinType,
                GeneratedFilterCondition generatedJoinCondition,
                @Nullable GeneratedFunction<FlatMapFunction<RowData, RowData>> generatedCalc,
                Node left,
                Node right,
                RowDataSerializer rowDataSerializerPassThroughCalc) {
            super(generatedCalc, rowDataSerializerPassThroughCalc);
            this.joinType = joinType;
            this.generatedJoinCondition = generatedJoinCondition;
            this.left = left;
            this.right = right;
        }

        @Override
        public void open(RuntimeContext runtimeContext, OpenContext openContext) throws Exception {
            super.open(runtimeContext, openContext);
            joinCondition =
                    generatedJoinCondition.newInstance(runtimeContext.getUserCodeClassLoader());
            FunctionUtils.setFunctionRuntimeContext(joinCondition, runtimeContext);
            FunctionUtils.openFunction(joinCondition, openContext);

            left.open(runtimeContext, openContext);
            right.open(runtimeContext, openContext);
        }

        @Override
        public void close() throws Exception {
            super.close();

            if (joinCondition != null) {
                FunctionUtils.closeFunction(joinCondition);
            }

            left.close();
            right.close();
        }

        @Override
        public Set<Integer> getAllInputOrdinals() {
            return Stream.concat(
                            left.getAllInputOrdinals().stream(),
                            right.getAllInputOrdinals().stream())
                    .collect(Collectors.toSet());
        }

        @Override
        public Node copy() {
            Node newLeft = left.copy();
            Node newRight = right.copy();
            return new JoinNode(
                    joinType,
                    generatedJoinCondition,
                    generatedCalc,
                    newLeft,
                    newRight,
                    rowDataSerializerPassThroughCalc);
        }
    }
}
