/*
 * 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.repair.autorepair;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.cassandra.config.DataStorageSpec;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.lifecycle.SSTableIntervalTree;
import org.apache.cassandra.db.lifecycle.SSTableSet;
import org.apache.cassandra.db.lifecycle.View;
import org.apache.cassandra.dht.Range;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.service.AutoRepairService;
import org.apache.cassandra.utils.concurrent.Refs;

import static org.apache.cassandra.repair.autorepair.AutoRepairUtils.split;

/**
 * The default implementation of {@link IAutoRepairTokenRangeSplitter} that attempts to:
 * <ol>
 *     <li>Create smaller, consistent repair times</li>
 *     <li>Minimize the impact on hosts</li>
 *     <li>Reduce overstreaming</li>
 *     <li>Reduce number of repairs</li>
 * </ol>
 * <p>
 * To achieve these goals, this implementation inspects SSTable metadata to estimate the bytes and number of partitions
 * within a range and splits it accordingly to bound the size of the token ranges used for repair assignments.
 * </p>
 * <p>
 * Refer to
 * <a href="https://cassandra.apache.org/doc/latest/cassandra/managing/operating/auto_repair.html#repair-token-range-splitter">Auto Repair documentation for this implementation</a>
 * for a more thorough breakdown of this implementation.
 * </p>
 * <p>
 * While this splitter has a lot of tuning parameters, the expectation is that the established default configuration
 * shall be sensible for all {@link org.apache.cassandra.repair.autorepair.AutoRepairConfig.RepairType}'s. The following
 * configuration parameters are offered.
 * </p>
 *
 * <p><b>Configuration parameters:</b></p>
 * <ul>
 *     <li><b>bytes_per_assignment</b> – Target size (in compressed bytes) for each repair. Throttles incremental repair
 *     and anticompaction per schedule after incremental repairs are enabled.</li>
 *
 *     <li><b>max_bytes_per_schedule</b> – Maximum data (in compressed bytes) to cover in a single schedule. Acts as a
 *     throttle for the repair cycle workload. Tune this up if writes are outpacing repair, or down if repairs are too
 *     disruptive. Alternatively, adjust {@code min_repair_interval}.</li>
 *
 *     <li><b>partitions_per_assignment</b> – Maximum number of partitions per repair assignment. Limits the number of
 *     partitions in Merkle tree leaves to prevent overstreaming.</li>
 *
 *     <li><b>max_tables_per_assignment</b> – Maximum number of tables to include in a single repair assignment.
 *     Especially useful for keyspaces with many tables. Prevents excessive batching of tables that exceed other
 *     parameters like {@code bytes_per_assignment} or {@code partitions_per_assignment}.</li>
 * </ul>
 */
public class RepairTokenRangeSplitter implements IAutoRepairTokenRangeSplitter
{
    private static final Logger logger = LoggerFactory.getLogger(RepairTokenRangeSplitter.class);

    // Default max bytes to 100TiB, which is much more readable than Long.MAX_VALUE
    private static final DataStorageSpec.LongBytesBound MAX_BYTES = new DataStorageSpec.LongBytesBound(102_400, DataStorageSpec.DataStorageUnit.GIBIBYTES);

    /**
     * The target bytes that should be included in a repair assignment
     */
    static final String BYTES_PER_ASSIGNMENT = "bytes_per_assignment";

    /**
     * Maximum number of partitions to include in a repair assignment
     */
    static final String PARTITIONS_PER_ASSIGNMENT = "partitions_per_assignment";

    /**
     * Maximum number of tables to include in a repair assignment if {@link AutoRepairConfig.Options#repair_by_keyspace}
     * is enabled
     */
    static final String MAX_TABLES_PER_ASSIGNMENT = "max_tables_per_assignment";

    /**
     * The maximum number of bytes to cover in an individual schedule
     */
    static final String MAX_BYTES_PER_SCHEDULE = "max_bytes_per_schedule";

    static final List<String> PARAMETERS = Arrays.asList(BYTES_PER_ASSIGNMENT, PARTITIONS_PER_ASSIGNMENT, MAX_TABLES_PER_ASSIGNMENT, MAX_BYTES_PER_SCHEDULE);

    private final AutoRepairConfig.RepairType repairType;

    private final Map<String, String> givenParameters = new HashMap<>();

    private DataStorageSpec.LongBytesBound bytesPerAssignment;
    private long partitionsPerAssignment;
    private int maxTablesPerAssignment;
    private DataStorageSpec.LongBytesBound maxBytesPerSchedule;

    /**
     * Established default for each {@link org.apache.cassandra.repair.autorepair.AutoRepairConfig.RepairType}, meant to
     * choose sensible defaults for each.
     * <p>
     * Defaults if not specified for the given repair type:
     * <li>
     *     <ul><b>bytes_per_assignment</b>: 50GiB</ul>
     *     <ul><b>partitions_per_assignment</b>: 1048576 (2^20)</ul>
     *     <ul><b>max_tables_per_assignment</b>: 64</ul>
     *     <ul><b>max_bytes_per_schedule</b>: 1000GiB</ul>
     * </li>
     * It's expected that these defaults should work well for everything except incremental, where we set
     * max_bytes_per_schedule to 100GiB. This should strike a good balance between the amount of data that will be
     * repaired during an initial migration to incremental repair and should move the entire repaired set from
     * unrepaired to repaired at steady state, assuming not more the 100GiB of data is written to a node per
     * min_repair_interval.
     */
    private static final Map<AutoRepairConfig.RepairType, RepairTypeDefaults> DEFAULTS_BY_REPAIR_TYPE = new EnumMap<AutoRepairConfig.RepairType, RepairTypeDefaults>(AutoRepairConfig.RepairType.class) {{
        put(AutoRepairConfig.RepairType.FULL, RepairTypeDefaults.builder(AutoRepairConfig.RepairType.FULL)
                                                                .build());
        // Restrict incremental repair to 100GiB max bytes per schedule to confine the amount of possible autocompaction.
        put(AutoRepairConfig.RepairType.INCREMENTAL, RepairTypeDefaults.builder(AutoRepairConfig.RepairType.INCREMENTAL)
                                                                       .withMaxBytesPerSchedule(new DataStorageSpec.LongBytesBound("100GiB"))
                                                                       .build());
        put(AutoRepairConfig.RepairType.PREVIEW_REPAIRED, RepairTypeDefaults.builder(AutoRepairConfig.RepairType.PREVIEW_REPAIRED)
                                                                            .build());
    }};

    public RepairTokenRangeSplitter(AutoRepairConfig.RepairType repairType, Map<String, String> parameters)
    {
        this.repairType = repairType;
        this.givenParameters.putAll(parameters);

        reinitParameters();
    }

    private void reinitParameters()
    {
        RepairTypeDefaults defaults = DEFAULTS_BY_REPAIR_TYPE.get(repairType);

        DataStorageSpec.LongBytesBound bytesPerAssignmentTmp = getPropertyOrDefault(BYTES_PER_ASSIGNMENT, DataStorageSpec.LongBytesBound::new, defaults.bytesPerAssignment);
        DataStorageSpec.LongBytesBound maxBytesPerScheduleTmp = getPropertyOrDefault(MAX_BYTES_PER_SCHEDULE, DataStorageSpec.LongBytesBound::new, defaults.maxBytesPerSchedule);

        // Validate that bytesPerAssignment <= maxBytesPerSchedule
        if (bytesPerAssignmentTmp.toBytes() > maxBytesPerScheduleTmp.toBytes())
        {
            throw new IllegalArgumentException(String.format("%s='%s' cannot be greater than %s='%s' for %s",
                                                             BYTES_PER_ASSIGNMENT,
                                                             bytesPerAssignmentTmp,
                                                             MAX_BYTES_PER_SCHEDULE,
                                                             maxBytesPerScheduleTmp,
                                                             repairType.getConfigName()));
        }

        bytesPerAssignment = bytesPerAssignmentTmp;
        maxBytesPerSchedule = maxBytesPerScheduleTmp;

        partitionsPerAssignment = getPropertyOrDefault(PARTITIONS_PER_ASSIGNMENT, Long::parseLong, defaults.partitionsPerAssignment);
        maxTablesPerAssignment = getPropertyOrDefault(MAX_TABLES_PER_ASSIGNMENT, Integer::parseInt, defaults.maxTablesPerAssignment);

        logger.info("Configured {}[{}] with {}={}, {}={}, {}={}, {}={}", RepairTokenRangeSplitter.class.getName(),
                    repairType.getConfigName(),
                    BYTES_PER_ASSIGNMENT, bytesPerAssignment,
                    PARTITIONS_PER_ASSIGNMENT, partitionsPerAssignment,
                    MAX_TABLES_PER_ASSIGNMENT, maxTablesPerAssignment,
                    MAX_BYTES_PER_SCHEDULE, maxBytesPerSchedule);
    }

    private <T> T getPropertyOrDefault(String propertyName, Function<String, T> mapper, T defaultValue)
    {
        return Optional.ofNullable(this.givenParameters.get(propertyName)).map(mapper).orElse(defaultValue);
    }

    @Override
    public Iterator<KeyspaceRepairAssignments> getRepairAssignments(boolean primaryRangeOnly, List<PrioritizedRepairPlan> repairPlans)
    {
        return new BytesBasedRepairAssignmentIterator(primaryRangeOnly, repairPlans);
    }

    /**
     * A custom {@link RepairAssignmentIterator} that confines the number of repair assignments to
     * <code>max_bytes_per_schedule</code>.
     */
    private class BytesBasedRepairAssignmentIterator extends RepairAssignmentIterator
    {

        private final boolean primaryRangeOnly;
        private long bytesSoFar = 0;

        BytesBasedRepairAssignmentIterator(boolean primaryRangeOnly, List<PrioritizedRepairPlan> repairPlans)
        {
            super(repairPlans);
            this.primaryRangeOnly = primaryRangeOnly;
        }

        @Override
        protected KeyspaceRepairAssignments next(int priority, KeyspaceRepairPlan repairPlan)
        {
            // short circuit if we've accumulated too many bytes by returning a KeyspaceRepairAssignments with
            // no assignments. We do this rather than returning false in hasNext() because we want to signal
            // to AutoRepair that a keyspace generated no assignments.
            if (bytesSoFar >= maxBytesPerSchedule.toBytes())
            {
                return new KeyspaceRepairAssignments(priority, repairPlan.getKeyspaceName(), Collections.emptyList());
            }

            List<Range<Token>> tokenRanges = AutoRepairUtils.getTokenRanges(primaryRangeOnly, repairPlan.getKeyspaceName());
            // shuffle token ranges to unbias selection of ranges
            Collections.shuffle(tokenRanges);
            List<SizedRepairAssignment> repairAssignments = new ArrayList<>();
            // Generate assignments for each range speparately
            for (Range<Token> tokenRange : tokenRanges)
            {
                repairAssignments.addAll(getRepairAssignmentsForKeyspace(repairType, repairPlan, tokenRange));
            }

            FilteredRepairAssignments filteredRepairAssignments = filterRepairAssignments(priority, repairPlan.getKeyspaceName(), repairAssignments, bytesSoFar);
            bytesSoFar = filteredRepairAssignments.newBytesSoFar;
            return new KeyspaceRepairAssignments(priority, repairPlan.getKeyspaceName(), filteredRepairAssignments.repairAssignments);
        }
    }

    @VisibleForTesting
    List<SizedRepairAssignment> getRepairAssignmentsForKeyspace(AutoRepairConfig.RepairType repairType, KeyspaceRepairPlan repairPlan, Range<Token> tokenRange)
    {
        List<SizedRepairAssignment> repairAssignments = new ArrayList<>();
        // this is used for batching minimal single assignment tables together
        List<SizedRepairAssignment> currentAssignments = new ArrayList<>();

        AutoRepairConfig config = AutoRepairService.instance.getAutoRepairConfig();

        // If we can repair by keyspace, sort the tables by size so can batch the smallest ones together
        boolean repairByKeyspace = config.getRepairByKeyspace(repairType);
        List<String> tablesToProcess = repairPlan.getTableNames();
        if (repairByKeyspace)
        {
            tablesToProcess = repairPlan.getTableNames().stream().sorted((t1, t2) -> {
                ColumnFamilyStore cfs1 = ColumnFamilyStore.getIfExists(repairPlan.getKeyspaceName(), t1);
                ColumnFamilyStore cfs2 = ColumnFamilyStore.getIfExists(repairPlan.getKeyspaceName(), t2);
                // If for whatever reason the CFS is not retrievable, we can assume it has been deleted, so give the
                // other cfs precedence.
                if (cfs1 == null)
                {
                    // cfs1 is lesser than because its null
                    return -1;
                }
                else if (cfs2 == null)
                {
                    // cfs1 is greather than because cfs2 is null
                    return 1;
                }
                return Long.compare(cfs1.metric.totalDiskSpaceUsed.getCount(), cfs2.metric.totalDiskSpaceUsed.getCount());
            }).collect(Collectors.toList());
        }

        for (String tableName : tablesToProcess)
        {
            List<SizedRepairAssignment> tableAssignments = getRepairAssignmentsForTable(repairPlan, tableName, tokenRange);

            if (tableAssignments.isEmpty())
                continue;

            // if not repairing by keyspace don't attempt to batch them with others.
            if (!repairByKeyspace)
            {
                repairAssignments.addAll(tableAssignments);
            }
            // If the table assignments are for the same token range, and we have room to add more tables to the current assignment
            else if (tableAssignments.size() == 1 &&
                     currentAssignments.size() < maxTablesPerAssignment &&
                     (currentAssignments.isEmpty() || currentAssignments.get(0).getTokenRange().equals(tableAssignments.get(0).getTokenRange())))
            {
                long currentAssignmentsBytes = getEstimatedBytes(currentAssignments);
                long tableAssignmentsBytes = getEstimatedBytes(tableAssignments);
                // only add assignments together if they don't exceed max bytes per schedule.
                if (currentAssignmentsBytes + tableAssignmentsBytes < maxBytesPerSchedule.toBytes()) {
                    currentAssignments.addAll(tableAssignments);
                }
                else
                {
                    // add table assignments by themselves
                    repairAssignments.addAll(tableAssignments);
                }
            }
            else
            {
                if (!currentAssignments.isEmpty())
                {
                    repairAssignments.add(merge(currentAssignments));
                    currentAssignments.clear();
                }
                repairAssignments.addAll(tableAssignments);
            }
        }

        if (!currentAssignments.isEmpty())
            repairAssignments.add(merge(currentAssignments));

        return repairAssignments;
    }

    /**
     * Given a repair type and map of sized-based repair assignments, confine them by <code>maxBytesPerSchedule</code>.
     *
     * @param repairAssignments the assignments to filter.
     * @param bytesSoFar        repair assignment bytes accumulated so far.
     * @return A list of repair assignments confined by <code>maxBytesPerSchedule</code>.
     */
    @VisibleForTesting
    FilteredRepairAssignments filterRepairAssignments(int priority, String keyspaceName, List<SizedRepairAssignment> repairAssignments, long bytesSoFar)
    {
        // Confine repair assignments by maxBytesPerSchedule.
        long bytesSoFarThisIteration = 0L;
        long bytesNotRepaired = 0L;
        int assignmentsNotRepaired = 0;
        int assignmentsToRepair = 0;
        int totalAssignments = 0;

        List<RepairAssignment> assignmentsToReturn = new ArrayList<>(repairAssignments.size());
        for (SizedRepairAssignment repairAssignment : repairAssignments)
        {
            totalAssignments++;
            // skip any repair assignments that would accumulate us past the maxBytesPerSchedule
            if (bytesSoFar + repairAssignment.getEstimatedBytes() > maxBytesPerSchedule.toBytes())
            {
                // log that repair assignment was skipped.
                bytesNotRepaired += repairAssignment.getEstimatedBytes();
                assignmentsNotRepaired++;
                logger.warn("Skipping {} because it would increase total repair bytes to {}",
                            repairAssignment,
                            getBytesOfMaxBytesPerSchedule(bytesSoFar + repairAssignment.getEstimatedBytes()));
            }
            else
            {
                bytesSoFar += repairAssignment.getEstimatedBytes();
                bytesSoFarThisIteration += repairAssignment.getEstimatedBytes();
                assignmentsToRepair++;
                logger.info("Adding {}, increasing repair bytes to {}",
                            repairAssignment,
                            getBytesOfMaxBytesPerSchedule(bytesSoFar));
                assignmentsToReturn.add(repairAssignment);
            }
        }

        String message = "Returning {} assignment(s) for priorityBucket {} and keyspace {}, totaling {} ({} overall)";
        if (assignmentsNotRepaired != 0)
        {
            message += ". Skipping {} of {} assignment(s), totaling {}";
            if (repairType != AutoRepairConfig.RepairType.INCREMENTAL)
            {
                message += ". The entire primary range will not be repaired this schedule. " +
                           "Consider increasing maxBytesPerSchedule, reducing node density or monitoring to ensure " +
                           "all ranges do get repaired within gc_grace_seconds";
                logger.warn(message, assignmentsToRepair, priority, keyspaceName,
                            FileUtils.stringifyFileSize(bytesSoFarThisIteration),
                            getBytesOfMaxBytesPerSchedule(bytesSoFar),
                            assignmentsNotRepaired, totalAssignments,
                            FileUtils.stringifyFileSize(bytesNotRepaired));
            }
            else
            {
                logger.info(message, assignmentsToRepair, priority, keyspaceName,
                            FileUtils.stringifyFileSize(bytesSoFarThisIteration),
                            getBytesOfMaxBytesPerSchedule(bytesSoFar),
                            assignmentsNotRepaired, totalAssignments,
                            FileUtils.stringifyFileSize(bytesNotRepaired));
            }
        }
        else
        {
            logger.info(message, assignmentsToRepair, priority, keyspaceName,
                        FileUtils.stringifyFileSize(bytesSoFarThisIteration),
                        getBytesOfMaxBytesPerSchedule(bytesSoFar));
        }

        return new FilteredRepairAssignments(assignmentsToReturn, bytesSoFar);
    }

    @VisibleForTesting
    static class FilteredRepairAssignments
    {
        final List<RepairAssignment> repairAssignments;
        final long newBytesSoFar;

        private FilteredRepairAssignments(List<RepairAssignment> repairAssignments, long newBytesSoFar)
        {
            this.repairAssignments = repairAssignments;
            this.newBytesSoFar = newBytesSoFar;
        }
    }

    private String getBytesOfMaxBytesPerSchedule(long bytes)
    {
        if (maxBytesPerSchedule.equals(MAX_BYTES))
            return FileUtils.stringifyFileSize(bytes);
        else
            return String.format("%s of %s", FileUtils.stringifyFileSize(bytes), maxBytesPerSchedule);
    }

    /**
     * @param repairAssignments The assignments to sum
     * @return The sum of {@link SizedRepairAssignment#getEstimatedBytes()} of all given
     * repairAssignments.
     */
    @VisibleForTesting
    protected static long getEstimatedBytes(List<SizedRepairAssignment> repairAssignments)
    {
        return repairAssignments
               .stream()
               .mapToLong(SizedRepairAssignment::getEstimatedBytes)
               .sum();
    }

    @VisibleForTesting
    static SizedRepairAssignment merge(List<SizedRepairAssignment> assignments)
    {
        if (assignments.isEmpty())
            throw new IllegalStateException("Cannot merge empty assignments");

        Set<String> mergedTableNames = new HashSet<>();
        Range<Token> referenceTokenRange = assignments.get(0).getTokenRange();
        String referenceKeyspaceName = assignments.get(0).getKeyspaceName();

        for (SizedRepairAssignment assignment : assignments)
        {
            // These checks _should_ be unnecessary but are here to ensure that the assignments are consistent
            if (!assignment.getTokenRange().equals(referenceTokenRange))
                throw new IllegalStateException("All assignments must have the same token range");
            if (!assignment.getKeyspaceName().equals(referenceKeyspaceName))
                throw new IllegalStateException("All assignments must have the same keyspace name");

            mergedTableNames.addAll(assignment.getTableNames());
        }

        long sizeForAssignment = getEstimatedBytes(assignments);
        return new SizedRepairAssignment(referenceTokenRange, referenceKeyspaceName, new ArrayList<>(mergedTableNames),
                                         "full primary range for " + mergedTableNames.size() + " tables", sizeForAssignment);
    }

    @VisibleForTesting
    protected List<SizedRepairAssignment> getRepairAssignmentsForTable(KeyspaceRepairPlan repairPlan, String tableName, Range<Token> tokenRange)
    {
        AutoRepairUtils.SizeEstimate sizeEstimate = repairPlan.getSizeEstimate(AutoRepairUtils.getKeyspaceTableName(repairPlan.getKeyspaceName(), tableName), tokenRange);
        if (sizeEstimate == null)
        {
            // Ideally, it should have been cached already inside the KeyspaceRepairPlan, but incase it was not,
            // then recalculating it. It is a bit expensive, but necessary for the repair
            logger.warn("The size estimate for {}.{} range {} was not pre-calculated, calculating on-demand",
                        repairPlan.getKeyspaceName(), tableName, tokenRange);
            sizeEstimate = AutoRepairUtils.getRangeSizeEstimate(repairType, repairPlan.getKeyspaceName(), tableName, tokenRange);
        }
        return getRepairAssignments(sizeEstimate);
    }

    private static void logSkippingTable(String keyspaceName, String tableName)
    {
        logger.warn("Could not resolve table data for {}.{} assuming it has since been deleted, skipping", keyspaceName, tableName);
    }

    @VisibleForTesting
    protected List<SizedRepairAssignment> getRepairAssignments(AutoRepairUtils.SizeEstimate estimate)
    {
        List<SizedRepairAssignment> repairAssignments = new ArrayList<>();

        // since its possible for us to hit maxBytesPerSchedule before seeing all ranges, shuffle so there is chance
        // at least of hitting all the ranges _eventually_ for the worst case scenarios
        int totalExpectedSubRanges = 0;
        if (estimate.sizeForRepair != 0)
        {
            boolean needsSplitting = estimate.sizeForRepair > bytesPerAssignment.toBytes() || estimate.partitions > partitionsPerAssignment;
            if (needsSplitting)
            {
                totalExpectedSubRanges += calculateNumberOfSplits(estimate);
            }
        }
        if (estimate.sizeForRepair == 0)
        {
            ColumnFamilyStore cfs = ColumnFamilyStore.getIfExists(estimate.keyspace, estimate.table);

            if (cfs == null)
            {
                logSkippingTable(estimate.keyspace, estimate.table);
                return Collections.emptyList();
            }

            long memtableSize = cfs.getTracker().getView().getCurrentMemtable().getLiveDataSize();
            if (memtableSize > 0L)
            {
                logger.debug("Included {}.{} range {}, had no unrepaired SSTables, but memtableSize={}, adding single repair assignment", estimate.keyspace, estimate.table, estimate.tokenRange, memtableSize);
                SizedRepairAssignment assignment = new SizedRepairAssignment(estimate.tokenRange, estimate.keyspace, Collections.singletonList(estimate.table), "full primary rangee for table with memtable only detected", memtableSize);
                repairAssignments.add(assignment);
            }
            else
            {
                logger.debug("Included {}.{} range {}, has no SSTables or memtable data, but adding single repair assignment for entire range in case writes were missed", estimate.keyspace, estimate.table, estimate.tokenRange);
                SizedRepairAssignment assignment = new SizedRepairAssignment(estimate.tokenRange, estimate.keyspace, Collections.singletonList(estimate.table), "full primary range for table with no data detected", 0L);
                repairAssignments.add(assignment);
            }
        }
        else
        {
            // Check if the estimate needs splitting based on the criteria
            boolean needsSplitting = estimate.sizeForRepair > bytesPerAssignment.toBytes() || estimate.partitions > partitionsPerAssignment;
            if (needsSplitting)
            {
                int numberOfSplits = calculateNumberOfSplits(estimate);
                long approximateBytesPerSplit = estimate.sizeForRepair / numberOfSplits;
                Collection<Range<Token>> subranges = split(estimate.tokenRange, numberOfSplits);
                for (Range<Token> subrange : subranges)
                {
                    SizedRepairAssignment assignment = new SizedRepairAssignment(subrange, estimate.keyspace, Collections.singletonList(estimate.table),
                                                                                 String.format("subrange %d of %d", repairAssignments.size() + 1, totalExpectedSubRanges),
                                                                                 approximateBytesPerSplit);
                    repairAssignments.add(assignment);
                }
            }
            else
            {
                // No splitting needed, repair the entire range as-is
                SizedRepairAssignment assignment = new SizedRepairAssignment(estimate.tokenRange, estimate.keyspace,
                                                                             Collections.singletonList(estimate.table),
                                                                             "full primary range for table", estimate.sizeForRepair);
                repairAssignments.add(assignment);
            }
        }
        return repairAssignments;
    }

    private int calculateNumberOfSplits(AutoRepairUtils.SizeEstimate estimate)
    {
        // Calculate the number of splits needed for size and partitions
        int splitsForSize = (int) Math.ceil((double) estimate.sizeForRepair / bytesPerAssignment.toBytes());
        int splitsForPartitions = (int) Math.ceil((double) estimate.partitions / partitionsPerAssignment);

        // Split the token range into subranges based on whichever (partitions, bytes) would generate the most splits.
        boolean splitBySize = splitsForSize > splitsForPartitions;
        int splits = splitBySize ? splitsForSize : splitsForPartitions;

        // calculate approximation for logging purposes
        long approximateBytesPerSplit = estimate.sizeForRepair / splits;
        long approximatePartitionsPerSplit = estimate.partitions / splits;

        logger.info("Splitting {}.{} for range {} into {} sub ranges by {} (splitsForSize={}, splitsForPartitions={}, " +
                    "approximateBytesInRange={}, approximatePartitionsInRange={}, " +
                    "approximateBytesPerSplit={}, approximatePartitionsPerSplit={})",
                    estimate.keyspace, estimate.table, estimate.tokenRange,
                    splits, splitBySize ? "size" : "partitions",
                    splitsForSize, splitsForPartitions,
                    FileUtils.stringifyFileSize(estimate.sizeForRepair), estimate.partitions,
                    FileUtils.stringifyFileSize(approximateBytesPerSplit), approximatePartitionsPerSplit
        );
        return splits;
    }

    @VisibleForTesting
    static Refs<SSTableReader> getSSTableReaderRefs(AutoRepairConfig.RepairType repairType, String keyspaceName, String tableName, Range<Token> tokenRange)
    {
        final ColumnFamilyStore cfs = ColumnFamilyStore.getIfExists(keyspaceName, tableName);
        if (cfs == null)
        {
            logSkippingTable(keyspaceName, tableName);
            return Refs.ref(Collections.<SSTableReader>emptyList());
        }

        Refs<SSTableReader> refs = null;
        while (refs == null)
        {
            Iterable<SSTableReader> sstables = cfs.getTracker().getView().select(SSTableSet.CANONICAL);
            SSTableIntervalTree tree = SSTableIntervalTree.buildSSTableIntervalTree(Lists.newArrayList(sstables));
            Range<PartitionPosition> r = Range.makeRowRange(tokenRange);
            List<SSTableReader> canonicalSSTables = View.sstablesInBounds(r.left, r.right, tree);
            if (repairType == AutoRepairConfig.RepairType.INCREMENTAL)
            {
                canonicalSSTables = canonicalSSTables.stream().filter((sstable) -> !sstable.isRepaired()).collect(Collectors.toList());
            }
            refs = Refs.tryRef(canonicalSSTables);
        }
        return refs;
    }

    @Override
    public void setParameter(String key, String value)
    {
        if (!PARAMETERS.contains(key))
        {
            throw new IllegalArgumentException("Unexpected parameter '" + key + "', must be one of " + PARAMETERS);
        }

        logger.info("Setting {} to {} for repair type {}", key, value, repairType);
        givenParameters.put(key, value);
        reinitParameters();
    }

    @Override
    public Map<String, String> getParameters()
    {
        final Map<String, String> parameters = new LinkedHashMap<>();
        for (String parameter : PARAMETERS)
        {
            // Use the parameter as provided if present.
            if (givenParameters.containsKey(parameter))
            {
                parameters.put(parameter, givenParameters.get(parameter));
                continue;
            }

            switch (parameter)
            {
                case BYTES_PER_ASSIGNMENT:
                    parameters.put(parameter, bytesPerAssignment.toString());
                    continue;
                case PARTITIONS_PER_ASSIGNMENT:
                    parameters.put(parameter, Long.toString(partitionsPerAssignment));
                    continue;
                case MAX_TABLES_PER_ASSIGNMENT:
                    parameters.put(parameter, Integer.toString(maxTablesPerAssignment));
                    continue;
                case MAX_BYTES_PER_SCHEDULE:
                    parameters.put(parameter, maxBytesPerSchedule.toString());
                    continue;
                default:
                    // not expected
                    parameters.put(parameter, "");
            }
        }
        return Collections.unmodifiableMap(parameters);
    }

    /**
     * Implementation of RepairAssignment that also assigns an estimation of bytes involved
     * in the repair.
     */
    @VisibleForTesting
    protected static class SizedRepairAssignment extends RepairAssignment
    {

        final String description;

        public SizedRepairAssignment(Range<Token> tokenRange, String keyspaceName, List<String> tableNames)
        {
            this(tokenRange, keyspaceName, tableNames, "", 0L);
        }

        public SizedRepairAssignment(Range<Token> tokenRange, String keyspaceName, List<String> tableNames,
                                     String description,
                                     long estimatedBytes)
        {
            super(tokenRange, keyspaceName, tableNames, estimatedBytes);
            this.description = description;
        }

        /**
         * @return Additional metadata about the repair assignment.
         */
        public String getDescription() {
            return description;
        }

        /**
         * Estimated bytes involved in the assignment. Typically Derived from {@link AutoRepairUtils.SizeEstimate#sizeForRepair}.
         *
         * @return estimated bytes involved in the assignment.
         */
        public long getEstimatedBytes()
        {
            return estimatedBytes;
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            if (!super.equals(o)) return false;
            SizedRepairAssignment that = (SizedRepairAssignment) o;
            return estimatedBytes == that.estimatedBytes && Objects.equals(description, that.description);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(super.hashCode(), description, estimatedBytes);
        }

        @Override
        public String toString()
        {
            return "SizedRepairAssignment{" +
                   "description='" + description + '\'' +
                   ", tokenRange=" + tokenRange +
                   ", keyspaceName='" + keyspaceName + '\'' +
                   ", tableNames=" + tableNames +
                   ", estimatedBytes=" + FileUtils.stringifyFileSize(estimatedBytes) +
                   '}';
        }
    }

    /**
     * Conveinence builder for establishing defaults by repair type.
     */
    protected static class RepairTypeDefaults
    {
        final AutoRepairConfig.RepairType repairType;
        final DataStorageSpec.LongBytesBound bytesPerAssignment;
        final long partitionsPerAssignment;
        final int maxTablesPerAssignment;
        final DataStorageSpec.LongBytesBound maxBytesPerSchedule;

        public RepairTypeDefaults(AutoRepairConfig.RepairType repairType,
                                  DataStorageSpec.LongBytesBound bytesPerAssignment,
                                  long partitionsPerAssignment,
                                  int maxTablesPerAssignment,
                                  DataStorageSpec.LongBytesBound maxBytesPerSchedule)
        {
            this.repairType = repairType;
            this.bytesPerAssignment = bytesPerAssignment;
            this.partitionsPerAssignment = partitionsPerAssignment;
            this.maxTablesPerAssignment = maxTablesPerAssignment;
            this.maxBytesPerSchedule = maxBytesPerSchedule;
        }

        static RepairTypeDefaultsBuilder builder(AutoRepairConfig.RepairType repairType)
        {
            return new RepairTypeDefaultsBuilder(repairType);
        }

        static class RepairTypeDefaultsBuilder
        {
            private final AutoRepairConfig.RepairType repairType;
            private DataStorageSpec.LongBytesBound bytesPerAssignment = new DataStorageSpec.LongBytesBound("50GiB");
            // Aims to target at most 1 partitions per leaf assuming a merkle tree of depth 20  (2^20 = 1,048,576)
            private long partitionsPerAssignment = 1_048_576;
            private int maxTablesPerAssignment = 64;
            private DataStorageSpec.LongBytesBound maxBytesPerSchedule = MAX_BYTES;

            private RepairTypeDefaultsBuilder(AutoRepairConfig.RepairType repairType)
            {
                this.repairType = repairType;
            }

            @SuppressWarnings("unused")
            public RepairTypeDefaultsBuilder withBytesPerAssignment(DataStorageSpec.LongBytesBound bytesPerAssignment)
            {
                this.bytesPerAssignment = bytesPerAssignment;
                return this;
            }

            @SuppressWarnings("unused")
            public RepairTypeDefaultsBuilder withPartitionsPerAssignment(long partitionsPerAssignment)
            {
                this.partitionsPerAssignment = partitionsPerAssignment;
                return this;
            }

            @SuppressWarnings("unused")
            public RepairTypeDefaultsBuilder withMaxTablesPerAssignment(int maxTablesPerAssignment)
            {
                this.maxTablesPerAssignment = maxTablesPerAssignment;
                return this;
            }

            public RepairTypeDefaultsBuilder withMaxBytesPerSchedule(DataStorageSpec.LongBytesBound maxBytesPerSchedule)
            {
                this.maxBytesPerSchedule = maxBytesPerSchedule;
                return this;
            }

            public RepairTokenRangeSplitter.RepairTypeDefaults build()
            {
                return new RepairTypeDefaults(repairType, bytesPerAssignment, partitionsPerAssignment, maxTablesPerAssignment, maxBytesPerSchedule);
            }
        }
    }
}
