/*
 * 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.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;

import org.apache.cassandra.SchemaLoader;
import org.apache.cassandra.config.DurationSpec;
import org.apache.cassandra.cql3.CQLTester;
import org.apache.cassandra.dht.BootStrapper;
import org.apache.cassandra.dht.Range;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.locator.TokenMetadata;
import org.apache.cassandra.schema.KeyspaceParams;
import org.apache.cassandra.service.AutoRepairService;
import org.apache.cassandra.service.StorageService;
import org.apache.cassandra.utils.FBUtilities;
import org.apache.cassandra.utils.Pair;

import static org.apache.cassandra.Util.setAutoRepairEnabled;
import static org.apache.cassandra.config.CassandraRelevantProperties.SYSTEM_DISTRIBUTED_DEFAULT_RF;
import static org.apache.cassandra.repair.autorepair.AutoRepairUtils.split;
import static org.apache.cassandra.repair.autorepair.FixedSplitTokenRangeSplitter.DEFAULT_NUMBER_OF_SUBRANGES;
import static org.apache.cassandra.repair.autorepair.FixedSplitTokenRangeSplitter.NUMBER_OF_SUBRANGES;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

/**
 * Helper class for {@link FixedSplitTokenRangeSplitterNoVNodesTest} and {@link FixedSplitTokenRangeSplitterVNodesTest}
 */
public class FixedSplitTokenRangeSplitterHelper extends CQLTester
{
    private static final String TABLE1 = "tbl1";
    private static final String TABLE2 = "tbl2";
    private static final String TABLE3 = "tbl3";
    public static final String KEYSPACE = "ks";
    public static final List<String> tables = Arrays.asList(TABLE1, TABLE2, TABLE3);
    public static final Map<String, Map<Range<Token>, AutoRepairUtils.SizeEstimate>> ksTablesEstimatedBytes = new HashMap<>();

    public static void setupClass(int numTokens) throws Exception
    {
        AutoRepair.SLEEP_IF_REPAIR_FINISHES_QUICKLY = new DurationSpec.IntSecondsBound("0s");
        SYSTEM_DISTRIBUTED_DEFAULT_RF.setInt(1);
        setAutoRepairEnabled(true);

        requireNetwork();
        AutoRepairUtils.setup();
        SchemaLoader.prepareServer();
        SchemaLoader.createKeyspace(KEYSPACE, KeyspaceParams.simple(1));

        TokenMetadata metadata = StorageService.instance.getTokenMetadata();
        metadata.updateNormalTokens(BootStrapper.getRandomTokens(metadata, numTokens), FBUtilities.getBroadcastAddressAndPort());

        Pair<Collection<Range<Token>>, Integer> tokensAndWrappedAroundCount = getTokenRangesAndTotalWrapAroundCount();
        int totalToken = numTokens + tokensAndWrappedAroundCount.right();
        long perTokenSizeTable1 = 512L / totalToken;
        long perTokenSizeTable2 = 1024L / totalToken;
        long perTokenSizeTable3 = 2048L / totalToken;
        for (Range<Token> tokenRange : tokensAndWrappedAroundCount.left)
        {
            ksTablesEstimatedBytes.put(AutoRepairUtils.getKeyspaceTableName(KEYSPACE, TABLE1), new HashMap<Range<Token>, AutoRepairUtils.SizeEstimate>()
            {{
                put(tokenRange, new AutoRepairUtils.SizeEstimate(AutoRepairConfig.RepairType.FULL, "", "", tokenRange, 0, perTokenSizeTable1, perTokenSizeTable1));
            }});
            ksTablesEstimatedBytes.put(AutoRepairUtils.getKeyspaceTableName(KEYSPACE, TABLE2), new HashMap<Range<Token>, AutoRepairUtils.SizeEstimate>()
            {{
                put(tokenRange, new AutoRepairUtils.SizeEstimate(AutoRepairConfig.RepairType.FULL, "", "", tokenRange, 0, perTokenSizeTable2, perTokenSizeTable2));
            }});
            ksTablesEstimatedBytes.put(AutoRepairUtils.getKeyspaceTableName(KEYSPACE, TABLE3), new HashMap<Range<Token>, AutoRepairUtils.SizeEstimate>()
            {{
                put(tokenRange, new AutoRepairUtils.SizeEstimate(AutoRepairConfig.RepairType.FULL, "", "", tokenRange, 0, perTokenSizeTable3, perTokenSizeTable3));
            }});
        }
    }

    public static void testTokenRangesSplitByTable(int numTokens, int numberOfSubRanges, AutoRepairConfig.RepairType repairType)
    {
        AutoRepairConfig config = AutoRepairService.instance.getAutoRepairConfig();
        config.setRepairByKeyspace(repairType, false);
        Pair<Collection<Range<Token>>, Integer> tokensAndWrappedAroundCount = getTokenRangesAndTotalWrapAroundCount();
        Collection<Range<Token>> tokens = tokensAndWrappedAroundCount.left();
        // For the test case, the tokens are allocated dynamically, so we do not know which token-ranges wrap around.
        // As a result, we need to adjust the token count on a need basis.
        numTokens += tokensAndWrappedAroundCount.right();
        assertEquals(numTokens, tokens.size());
        List<Range<Token>> expectedToken = new ArrayList<>();
        int numberOfSplits = Math.max(1, numberOfSubRanges / tokens.size());
        for (int i = 0; i < tables.size(); i++)
        {
            for (Range<Token> token : tokens)
            {
                expectedToken.addAll(split(token, numberOfSplits));
            }
        }

        Iterator<KeyspaceRepairAssignments> keyspaceAssignments =
        new FixedSplitTokenRangeSplitter(repairType, Collections.singletonMap(NUMBER_OF_SUBRANGES, Integer.toString(numberOfSubRanges)))
        .getRepairAssignments(config.getRepairPrimaryTokenRangeOnly(repairType), getPlan(repairType));

        // should be only 1 entry for the keyspace.
        assertTrue(keyspaceAssignments.hasNext());
        KeyspaceRepairAssignments keyspaceRepairAssignment = keyspaceAssignments.next();
        assertFalse(keyspaceAssignments.hasNext());

        List<RepairAssignment> assignments = keyspaceRepairAssignment.getRepairAssignments();
        assertEquals(numTokens * numberOfSplits * tables.size(), assignments.size());
        assertEquals(expectedToken.size(), assignments.size());

        int assignmentsPerTable = numTokens * numberOfSplits;
        for (int i = 0; i < tables.size(); i++)
        {
            List<RepairAssignment> assignmentForATable = new ArrayList<>();
            List<Range<Token>> expectedTokensForATable = new ArrayList<>();
            for (int j = 0; j < assignmentsPerTable; j++)
            {
                long expectedBytes = ksTablesEstimatedBytes.get(AutoRepairUtils.getKeyspaceTableName(KEYSPACE, tables.get(i))).values().stream().mapToLong(sizeEstimate -> sizeEstimate.sizeForRepair).sum() / numberOfSplits;
                int theTableAssignmentIdx = i * assignmentsPerTable + j;
                assertEquals(expectedBytes, assignments.get(theTableAssignmentIdx).estimatedBytes);
                assertEquals(Collections.singletonList(tables.get(i)), assignments.get(theTableAssignmentIdx).getTableNames());
                assignmentForATable.add(assignments.get(theTableAssignmentIdx));
                expectedTokensForATable.add(expectedToken.get(theTableAssignmentIdx));
            }
            compare(numTokens, numberOfSplits, expectedTokensForATable, assignmentForATable);
        }
    }

    public static void testTokenRangesSplitByKeyspace(int numTokens, int numberOfSubRanges, AutoRepairConfig.RepairType repairType)
    {
        AutoRepairConfig config = AutoRepairService.instance.getAutoRepairConfig();
        config.setRepairByKeyspace(repairType, true);
        Pair<Collection<Range<Token>>, Integer> tokensAndWrappedRanges = getTokenRangesAndTotalWrapAroundCount();
        Collection<Range<Token>> tokens = tokensAndWrappedRanges.left();
        // For the test case, the tokens are allocated dynamically, so we do not know which token-ranges wrap around.
        // As a result, we need to adjust the token count on a need basis.
        numTokens += tokensAndWrappedRanges.right();
        assertEquals(numTokens, tokens.size());
        int numberOfSplits = Math.max(1, numberOfSubRanges / tokens.size());
        List<Range<Token>> expectedToken = new ArrayList<>();
        for (Range<Token> range : tokens)
        {
            expectedToken.addAll(AutoRepairUtils.split(range, numberOfSplits));
        }

        Iterator<KeyspaceRepairAssignments> keyspaceAssignments =
        new FixedSplitTokenRangeSplitter(repairType, Collections.singletonMap(NUMBER_OF_SUBRANGES, Integer.toString(numberOfSubRanges)))
        .getRepairAssignments(config.getRepairPrimaryTokenRangeOnly(repairType), getPlan(repairType));

        // should be only 1 entry for the keyspace.
        assertTrue(keyspaceAssignments.hasNext());
        KeyspaceRepairAssignments keyspace = keyspaceAssignments.next();
        assertFalse(keyspaceAssignments.hasNext());

        List<RepairAssignment> assignments = keyspace.getRepairAssignments();
        assertNotNull(assignments);

        assertEquals(numTokens * numberOfSplits, assignments.size());
        assertEquals(expectedToken.size(), assignments.size());

        compare(numTokens, numberOfSplits, expectedToken, assignments);

        for (int i = 0; i < assignments.size(); i++)
        {
            assertEquals(assignments.get(i).estimatedBytes,
                         ksTablesEstimatedBytes.values().stream()
                                               .flatMap(tableMap -> tableMap.values().stream())
                                               .mapToLong(sizeEstimate -> sizeEstimate.sizeForRepair)
                                               .sum() / numberOfSplits);
        }
    }

    public static void testTokenRangesWithDefaultSplit(int numTokens, AutoRepairConfig.RepairType repairType)
    {
        testTokenRangesSplitByKeyspace(numTokens, DEFAULT_NUMBER_OF_SUBRANGES, repairType);
    }

    private static void compare(int numTokens, int numberOfSplits, List<Range<Token>> expectedToken, List<RepairAssignment> assignments)
    {
        assertEquals(expectedToken.size(), assignments.size());
        Set<Range<Token>> a = new TreeSet<>();
        Set<Range<Token>> b = new TreeSet<>();
        for (int i = 0; i < numTokens * numberOfSplits; i++)
        {
            a.add(expectedToken.get(i));
            b.add(assignments.get(i).getTokenRange());
        }
        assertEquals(a, b);
    }

    private static Pair<Collection<Range<Token>>, Integer> getTokenRangesAndTotalWrapAroundCount()
    {
        int wrappedRanges = 0;
        Collection<Range<Token>> ranges = StorageService.instance.getPrimaryRangesForEndpoint(KEYSPACE, FBUtilities.getBroadcastAddressAndPort());
        Collection<Range<Token>> tokens = new ArrayList<>();
        for (Range<Token> wrappedRange : ranges)
        {
            if (wrappedRange.isWrapAround())
            {
                wrappedRanges++;
            }
            tokens.addAll(wrappedRange.unwrap());
        }
        return Pair.create(tokens, wrappedRanges);
    }

    private static List<PrioritizedRepairPlan> getPlan(AutoRepairConfig.RepairType repairType)
    {
        AutoRepairConfig config = AutoRepairService.instance.getAutoRepairConfig();
        List<PrioritizedRepairPlan> plan = PrioritizedRepairPlan.build(new HashMap<String, List<String>>()
                                                                       {{
                                                                           put(KEYSPACE, tables);
                                                                       }}, repairType, (l) -> {
                                                                       },
                                                                       config.getRepairPrimaryTokenRangeOnly(repairType));
        assertEquals(1, plan.size());
        assertEquals(1, plan.get(0).getKeyspaceRepairPlans().size());
        plan.get(0).getKeyspaceRepairPlans().get(0).ksTablesEstimatedBytes = ksTablesEstimatedBytes;
        return plan;
    }
}
