/*
 * 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.connectors.kudu.table;

import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.connectors.kudu.connector.KuduFilterInfo;
import org.apache.flink.connectors.kudu.connector.KuduTableInfo;
import org.apache.flink.connectors.kudu.connector.convertor.RowResultRowConvertor;
import org.apache.flink.connectors.kudu.connector.reader.KuduReaderConfig;
import org.apache.flink.connectors.kudu.format.KuduRowInputFormat;
import org.apache.flink.shaded.guava30.com.google.common.collect.Lists;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.TableSchema;
import org.apache.flink.table.expressions.Expression;
import org.apache.flink.table.sources.*;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.utils.TypeConversions;
import org.apache.flink.types.Row;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import java.util.*;

import static org.apache.flink.connectors.kudu.table.utils.KuduTableUtils.toKuduFilterInfo;

public class KuduTableSource implements StreamTableSource<Row>,
        LimitableTableSource<Row>, ProjectableTableSource<Row>, FilterableTableSource<Row> {


    private static final Logger LOG = LoggerFactory.getLogger(KuduTableSource.class);

    private final KuduReaderConfig.Builder configBuilder;
    private final KuduTableInfo tableInfo;
    private final TableSchema flinkSchema;
    private final String[] projectedFields;
    // predicate expression to apply
    @Nullable
    private final List<KuduFilterInfo> predicates;
    private boolean isFilterPushedDown;

    private KuduRowInputFormat kuduRowInputFormat;

    public KuduTableSource(KuduReaderConfig.Builder configBuilder, KuduTableInfo tableInfo,
                           TableSchema flinkSchema, List<KuduFilterInfo> predicates, String[] projectedFields) {
        this.configBuilder = configBuilder;
        this.tableInfo = tableInfo;
        this.flinkSchema = flinkSchema;
        this.predicates = predicates;
        this.projectedFields = projectedFields;
        if (predicates != null && predicates.size() != 0) {
            this.isFilterPushedDown = true;
        }
        this.kuduRowInputFormat = new KuduRowInputFormat(configBuilder.build(), new RowResultRowConvertor(), tableInfo,
                predicates == null ? Collections.emptyList() : predicates,
                projectedFields == null ? null : Lists.newArrayList(projectedFields));
    }

    @Override
    public boolean isBounded() {
        return true;
    }

    @Override
    public DataStream<Row> getDataStream(StreamExecutionEnvironment env) {
        KuduRowInputFormat inputFormat = new KuduRowInputFormat(configBuilder.build(), new RowResultRowConvertor(),
                tableInfo,
                predicates == null ? Collections.emptyList() : predicates,
                projectedFields == null ? null : Lists.newArrayList(projectedFields));
        return env.createInput(inputFormat,
                        (TypeInformation<Row>) TypeConversions.fromDataTypeToLegacyInfo(getProducedDataType()))
                .name(explainSource());
    }

    @Override
    public TableSchema getTableSchema() {
        return flinkSchema;
    }

    @Override
    public boolean isFilterPushedDown() {
        return this.isFilterPushedDown;
    }

    @Override
    public DataType getProducedDataType() {
        if (projectedFields == null) {
            return flinkSchema.toRowDataType();
        } else {
            DataTypes.Field[] fields = new DataTypes.Field[projectedFields.length];
            for (int i = 0; i < fields.length; i++) {
                String fieldName = projectedFields[i];
                fields[i] = DataTypes.FIELD(
                        fieldName,
                        flinkSchema
                                .getTableColumn(fieldName)
                                .get()
                                .getType()
                );
            }
            return DataTypes.ROW(fields);
        }
    }

    @Override
    public boolean isLimitPushedDown() {
        return true;
    }

    @Override
    public TableSource<Row> applyLimit(long l) {
        return new KuduTableSource(configBuilder.setRowLimit((int) l), tableInfo, flinkSchema,
                predicates, projectedFields);
    }

    @Override
    public TableSource<Row> projectFields(int[] ints) {
        String[] fieldNames = new String[ints.length];
        RowType producedDataType = (RowType) getProducedDataType().getLogicalType();
        List<String> prevFieldNames = producedDataType.getFieldNames();
        for (int i = 0; i < ints.length; i++) {
            fieldNames[i] = prevFieldNames.get(ints[i]);
        }
        return new KuduTableSource(configBuilder, tableInfo, flinkSchema, predicates, fieldNames);
    }

    @Override
    public TableSource<Row> applyPredicate(List<Expression> predicates) {
        List<KuduFilterInfo> kuduPredicates = new ArrayList<>();
        ListIterator<Expression> predicatesIter = predicates.listIterator();
        while (predicatesIter.hasNext()) {
            Expression predicate = predicatesIter.next();
            Optional<KuduFilterInfo> kuduPred = toKuduFilterInfo(predicate);
            if (kuduPred != null && kuduPred.isPresent()) {
                LOG.debug("Predicate [{}] converted into KuduFilterInfo and pushed into " +
                        "KuduTable [{}].", predicate, tableInfo.getName());
                kuduPredicates.add(kuduPred.get());
                predicatesIter.remove();
            } else {
                LOG.debug("Predicate [{}] could not be pushed into KuduFilterInfo for KuduTable [{}].",
                        predicate, tableInfo.getName());
            }
        }
        return new KuduTableSource(configBuilder, tableInfo, flinkSchema, kuduPredicates, projectedFields);
    }

    @Override
    public String explainSource() {
        return "KuduTableSource[schema=" + Arrays.toString(getTableSchema().getFieldNames()) +
                ", filter=" + predicateString() +
                (projectedFields != null ? ", projectFields=" + Arrays.toString(projectedFields) + "]" : "]");
    }

    private String predicateString() {
        if (predicates == null || predicates.size() == 0) {
            return "No predicates push down";
        } else {
            return "AND(" + predicates + ")";
        }
    }
}
