/*
 * 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.cassandra.spark.sparksql;

import java.util.Arrays;
import java.util.List;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.apache.cassandra.spark.data.CqlField;
import org.apache.cassandra.spark.data.CqlTable;
import org.apache.cassandra.spark.data.DataLayer;
import org.apache.cassandra.spark.sparksql.filters.PartitionKeyFilter;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

/**
 * Wrapper iterator around SparkCellIterator to normalize cells into Spark SQL rows
 *
 * @param <T> type of row returned by Iterator.
 */
abstract class AbstractSparkRowIterator<T> extends RowIterator<T>
{
    AbstractSparkRowIterator(int partitionId,
                             @NotNull DataLayer dataLayer,
                             @Nullable StructType requiredSchema,
                             @NotNull List<PartitionKeyFilter> partitionKeyFilters,
                             Function<RowBuilder<T>, RowBuilder<T>> decorator)
    {
        super(
        buildCellIterator(partitionId, dataLayer.cqlTable(), requiredSchema, dataLayer, partitionKeyFilters),
        dataLayer.stats(),
        requiredSchema == null ? null : requiredSchema.fieldNames(),
        decorator
        );
    }

    protected static CellIterator buildCellIterator(int partitionId,
                                                    CqlTable cqlTable,
                                                    @Nullable StructType requiredSchema,
                                                    @NotNull DataLayer dataLayer,
                                                    @NotNull List<PartitionKeyFilter> partitionKeyFilters)
    {
        StructType columnFilter = useColumnFilter(requiredSchema, cqlTable) ? requiredSchema : null;
        return new SparkCellIterator(partitionId, dataLayer, columnFilter, partitionKeyFilters);
    }

    private static boolean useColumnFilter(@Nullable StructType requiredSchema, CqlTable cqlTable)
    {
        if (requiredSchema == null)
        {
            return false;
        }
        // Only use column filter if it excludes any of the CqlTable fields
        Set<String> requiredFields = Arrays.stream(requiredSchema.fields()).map(StructField::name).collect(Collectors.toSet());
        return cqlTable.fields().stream()
                       .map(CqlField::name)
                       .anyMatch(field -> !requiredFields.contains(field));
    }

    /**
     * Maps the Object[] valueArray generated by the RowIterator to the expected type
     *
     * @param valueArray
     * @return a value of type `T`
     */
    public abstract T rowBuilder(Object[] valueArray);

    @SuppressWarnings("DataFlowIssue") // Null check performed within PartialRowBuilder constructor
    @Override
    public PartialRowBuilder<T> newPartialBuilder()
    {
        return new PartialRowBuilder<>(requiredColumns, it.cqlTable(), it.hasProjectedValueColumns(), this::rowBuilder);
    }

    @Override
    public FullRowBuilder<T> newFullRowBuilder()
    {
        return new FullRowBuilder<>(it.cqlTable(), it.hasProjectedValueColumns(), this::rowBuilder);
    }
}
