/*
 * 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.harry.dsl;

import java.util.ArrayList;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.function.Predicate;

import org.apache.cassandra.harry.MagicConstants;
import org.apache.cassandra.harry.Relations;
import org.apache.cassandra.harry.SchemaSpec;
import org.apache.cassandra.harry.gen.EntropySource;
import org.apache.cassandra.harry.gen.Generator;
import org.apache.cassandra.harry.gen.Generators;
import org.apache.cassandra.harry.util.BitSet;

import static org.apache.cassandra.harry.Relations.RelationKind.EQ;
import static org.apache.cassandra.harry.Relations.RelationKind.GT;
import static org.apache.cassandra.harry.Relations.RelationKind.GTE;
import static org.apache.cassandra.harry.Relations.RelationKind.LT;

/**
 * Things that seemed like a good idea, but ultimately were not a good fit for the HistoryBuilder API
 */
public class HistoryBuilderHelper
{
    /**
     * Perform a random insert to any row
     */
    public static void insertRandomData(SchemaSpec schema, Generator<Integer> pkGen, Generator<Integer> ckGen, EntropySource rng, HistoryBuilder history)
    {
        insertRandomData(schema, pkGen.generate(rng), ckGen.generate(rng), rng, history);
    }

    public static void insertRandomData(SchemaSpec schema, int partitionIdx, int rowIdx, EntropySource rng, HistoryBuilder history)
    {
        int[] vIdxs = new int[schema.regularColumns.size()];
        for (int i = 0; i < schema.regularColumns.size(); i++)
            vIdxs[i] = rng.nextInt(history.valueGenerators().regularPopulation(i));
        int[] sIdxs = new int[schema.staticColumns.size()];
        for (int i = 0; i < schema.staticColumns.size(); i++)
            sIdxs[i] = rng.nextInt(history.valueGenerators().staticPopulation(i));
        history.insert(partitionIdx, rowIdx, vIdxs, sIdxs);
    }

    public static void insertRandomData(SchemaSpec schema, int pkIdx, EntropySource rng, HistoryBuilder history)
    {
        insertRandomData(schema,
                         pkIdx,
                         rng.nextInt(0, history.valueGenerators().ckPopulation()),
                         rng,
                         0,
                         history);
    }

    public static void insertRandomData(SchemaSpec schema, int partitionIdx, int rowIdx, EntropySource rng, double chanceOfUnset, HistoryBuilder history)
    {
        int[] vIdxs = new int[schema.regularColumns.size()];
        for (int i = 0; i < schema.regularColumns.size(); i++)
            vIdxs[i] = rng.nextDouble() <= chanceOfUnset ? MagicConstants.UNSET_IDX : rng.nextInt(history.valueGenerators().regularPopulation(i));
        int[] sIdxs = new int[schema.staticColumns.size()];
        for (int i = 0; i < schema.staticColumns.size(); i++)
            sIdxs[i] = rng.nextDouble() <= chanceOfUnset ? MagicConstants.UNSET_IDX : rng.nextInt(history.valueGenerators().staticPopulation(i));
        history.insert(partitionIdx, rowIdx, vIdxs, sIdxs);
    }

    public static void deleteRandomColumns(SchemaSpec schema, int partitionIdx, int rowIdx, EntropySource rng, SingleOperationBuilder history)
    {
        Generator<BitSet> regularMask = Generators.bitSet(schema.regularColumns.size());
        Generator<BitSet> staticMask = Generators.bitSet(schema.staticColumns.size());

        history.deleteColumns(partitionIdx,
                              rowIdx,
                              regularMask.generate(rng),
                              staticMask.generate(rng));
    }

    private static final Generator<Relations.RelationKind> relationKindGen = Generators.pick(LT, GT, EQ);
    private static final Set<Relations.RelationKind> lowBoundRelations = Set.of(GT, GTE, EQ);

    public static List<SingleOperationBuilder.IdxRelation> generateValueRelations(EntropySource rng, int numColumns, Function<Integer, Integer> population)
    {
        return generateValueRelations(rng, numColumns, population, c -> false);
    }

    /**
     * Generates random relations for regular and static columns for FILTERING and SAI queries.
     *
     * Will generate at most 2 relations per column:
     *   * generates a random relation
     *     * if this relation is EQ, that's the only relation that will lock this column
     *     * if relation is GT, next bound, if generated, will be LT
     *     * if relation is LT, next bound, if generated, will be GT
     *
     * @param rng - random number generator
     * @param numColumns - number of columns in the generated set of relationships
     * @param population - expected population / number of possible values for a given column
     * @return a list of relations
     */
    public static List<SingleOperationBuilder.IdxRelation> generateValueRelations(EntropySource rng, int numColumns, Function<Integer, Integer> population, Predicate<Integer> eqOnlyColumns)
    {
        List<SingleOperationBuilder.IdxRelation> relations = new ArrayList<>();
        Map<Integer, Set<Relations.RelationKind>> kindsMap = new HashMap<>();
        int remainingColumns = numColumns;
        while (remainingColumns > 0)
        {
            int column = rng.nextInt(numColumns);
            Set<Relations.RelationKind> kinds = kindsMap.computeIfAbsent(column, c -> EnumSet.noneOf(Relations.RelationKind.class));
            if (kinds.size() > 1 || kinds.contains(EQ))
                continue;
            Relations.RelationKind kind;
            
            if (eqOnlyColumns.test(column))
            {
                kind = EQ;
            }
            else if (kinds.size() == 1)
            {
                if (kinds.contains(LT)) kind = GT;
                else kind = LT;
                remainingColumns--;
            }
            else
                // TODO: weights per relation?
                kind = relationKindGen.generate(rng);

            if (kind == EQ)
                remainingColumns--;

            kinds.add(kind);

            int regularIdx = rng.nextInt(population.apply(column));
            relations.add(new SingleOperationBuilder.IdxRelation(kind, regularIdx, column));
            if (rng.nextBoolean())
                break;
        }
        return relations;
    }

    public static List<SingleOperationBuilder.IdxRelation> generateClusteringRelations(EntropySource rng, int numColumns, Generator<Integer> ckGen)
    {
        return generateClusteringRelations(rng, numColumns, ckGen, Collections.emptySet());
    }

    /**
     * Generates random relations for regular and static columns for FILTERING and SAI queries.
     *
     * Will generate at most 2 relations per column. Low bound will always use values from low bound clustering,
     * high bound will always use values from high bound.
     *
     * @param rng - random number generator
     * @param numColumns - number of columns in the generated set of relationships
     * @return a list of relations
     */
    public static List<SingleOperationBuilder.IdxRelation> generateClusteringRelations(EntropySource rng, int numColumns, Generator<Integer> ckGen, Set<Integer> eqOnlyColumns)
    {
        List<SingleOperationBuilder.IdxRelation> relations = new ArrayList<>();
        Map<Integer, Set<Relations.RelationKind>> kindsMap = new HashMap<>();
        int remainingColumns = numColumns;
        int lowBoundIdx = ckGen.generate(rng);
        int highBoundIdx = ckGen.generate(rng);
        while (remainingColumns > 0)
        {
            int column = rng.nextInt(numColumns);
            Set<Relations.RelationKind> kinds = kindsMap.computeIfAbsent(column, c -> EnumSet.noneOf(Relations.RelationKind.class));
            if (kinds.size() > 1 || kinds.contains(EQ))
                continue;
            Relations.RelationKind kind;

            if (eqOnlyColumns.contains(column))
            {
                kind = EQ;
            }
            else if (kinds.size() == 1)
            {
                if (kinds.contains(LT)) kind = GT;
                else kind = LT;
                remainingColumns--;
            }
            else
                kind = relationKindGen.generate(rng);

            if (kind == EQ)
                remainingColumns--;

            kinds.add(kind);

            relations.add(new SingleOperationBuilder.IdxRelation(kind, lowBoundRelations.contains(kind) ? lowBoundIdx : highBoundIdx, column));
            if (rng.nextBoolean())
                break;
        }
        return relations;
    }
}
