/*
 * Basic tests for moveRange.
 *
 * @tags: [
 *    assumes_balancer_off
 * ]
 */
import {FeatureFlagUtil} from "jstests/libs/feature_flag_util.js";
import {ShardingTest} from "jstests/libs/shardingtest.js";
import {chunkBoundsUtil} from "jstests/sharding/libs/chunk_bounds_util.js";
import {findChunksUtil} from "jstests/sharding/libs/find_chunks_util.js";

const st = new ShardingTest({mongos: 1, shards: 2, chunkSize: 1});
const kDbName = "db";

const mongos = st.s0;
const shard0 = st.shard0.shardName;
const shard1 = st.shard1.shardName;

assert.commandWorked(mongos.adminCommand({enableSharding: kDbName, primaryShard: shard0}));

function getRandomShardKeyValue(ns, skPattern, filter) {
    const isHashedShardKey = Object.values(skPattern).includes("hashed");
    const coll = mongos.getCollection(ns);

    // Get a random document from the collection
    let doc = coll.aggregate([{$match: filter}, {$sample: {size: 1}}]).next();

    // Delete fields not making part of the shard key
    for (let key in doc) {
        if (!(key in skPattern)) {
            delete doc[key];
        }
    }

    if (isHashedShardKey) {
        doc.a = convertShardKeyToHashed(doc.a);
    }

    return doc;
}

// Tests for moveRange that will call move&split. We call moveRange with minBound on the positive
// values and moveRange with maxBound on negative values to ensure that the chunk chosen is big
// enough to be split by moveRange.
function testMoveRangeWithBigChunk(mongos, ns, skPattern, minBound) {
    // Get a random existing shard key value, `moveRange` will be called on the owning chunk
    let filter = minBound ? {a: {$gte: 0}} : {a: {$lt: 0}};
    let randomSK = getRandomShardKeyValue(ns, skPattern, filter);

    // Get bounds and shard of the chunk owning `randomSK`
    const chunksBefore = findChunksUtil.findChunksByNs(mongos.getDB("config"), ns).toArray();
    const shardChunkBounds = chunkBoundsUtil.findShardChunkBounds(chunksBefore);
    const {shard, bounds} = chunkBoundsUtil.findShardAndChunkBoundsForShardKey(st, shardChunkBounds, randomSK);

    const donor = shard.shardName;
    const recipient = donor == shard0 ? shard1 : shard0;

    // Count chunks belonging to donor and recipient shards BEFORE moveRange
    const nChunksOnDonorBefore = chunksBefore.filter((chunk) => chunk.shard == donor).length;
    const nChunksOnRecipientBefore = chunksBefore.filter((chunk) => chunk.shard == recipient).length;

    if (minBound) {
        assert.commandWorked(mongos.adminCommand({moveRange: ns, min: randomSK, toShard: recipient}));
    } else {
        assert.commandWorked(mongos.adminCommand({moveRange: ns, max: randomSK, toShard: recipient}));
    }

    // Count chunks belonging to donor and recipient shards AFTER moveRange
    const chunksAfter = findChunksUtil.findChunksByNs(mongos.getDB("config"), ns).toArray();
    const nChunksOnDonorAfter = chunksAfter.filter((chunk) => chunk.shard == donor).length;
    const nChunksOnRecipientAfter = chunksAfter.filter((chunk) => chunk.shard == recipient).length;

    let nExpectedChunksOnRecipientAfter = nChunksOnRecipientBefore + 1;
    // For moveRange with a maxBound, the number of chunks on recipient and donor doesn't change
    // if a shardKey that is a lower-bound of a pre-existing chunk was selected.
    if (!minBound && chunkBoundsUtil.eq(bounds[0], randomSK)) {
        nExpectedChunksOnRecipientAfter = nChunksOnRecipientBefore;
    }

    assert.eq(
        nExpectedChunksOnRecipientAfter,
        nChunksOnRecipientAfter,
        "The number of chunks on the recipient shard did not increase following a moveRange",
    );
    assert(
        nChunksOnDonorAfter == nChunksOnDonorBefore || nChunksOnDonorAfter == nChunksOnDonorBefore + 1,
        "Unexpected number of chunks on the donor shard after triggering a split + move",
    );
}

function test(collName, skPattern, splitpoint) {
    const ns = kDbName + "." + collName;

    assert.commandWorked(mongos.adminCommand({shardCollection: ns, key: skPattern}));

    let aChunk = findChunksUtil.findOneChunkByNs(mongos.getDB("config"), ns, {shard: shard0});
    assert(aChunk);

    jsTest.log("Testing invalid commands");
    // Fail if one of the bounds is not a valid shard key
    assert.commandFailedWithCode(
        mongos.adminCommand({moveRange: ns, min: aChunk.min, max: {invalidShardKey: 10}, toShard: shard1}),
        ErrorCodes.InvalidOptions,
    );

    // Fail if the `to` shard does not exists
    assert.commandFailedWithCode(
        mongos.adminCommand({moveRange: ns, min: aChunk.min, max: aChunk.max, toShard: "WrongShard"}),
        ErrorCodes.ShardNotFound,
    );

    // Test that `moveRange` with min & max bounds works
    jsTest.log("Testing moveRange with both bounds");
    assert.commandWorked(mongos.adminCommand({moveRange: ns, min: aChunk.min, max: aChunk.max, toShard: shard1}));

    assert.eq(0, mongos.getDB("config").chunks.countDocuments({_id: aChunk._id, shard: shard0}));
    assert.eq(1, mongos.getDB("config").chunks.countDocuments({_id: aChunk._id, shard: shard1}));

    // Test that `moveRange` only with min bound works (translates to `moveChunk` because chunk too
    // small to be split)
    jsTest.log("Testing moveRange with only min bound");
    assert.commandWorked(mongos.adminCommand({moveRange: ns, min: aChunk.min, toShard: shard0}));

    assert.eq(1, mongos.getDB("config").chunks.countDocuments({_id: aChunk._id, shard: shard0}));
    assert.eq(0, mongos.getDB("config").chunks.countDocuments({_id: aChunk._id, shard: shard1}));

    // Test that `moveRange` only with max bound works (translates to `moveChunk` because chunk too
    // small to be split)
    jsTest.log("Testing moveRange with only max bound");
    assert.commandWorked(mongos.adminCommand({moveRange: ns, max: aChunk.max, toShard: shard1}));

    assert.eq(0, mongos.getDB("config").chunks.countDocuments({_id: aChunk._id, shard: shard0}));
    assert.eq(1, mongos.getDB("config").chunks.countDocuments({_id: aChunk._id, shard: shard1}));

    // Insert 10MB >0 and <0 in order to create multiple big chunks (chunkSize is set to 1MB)
    jsTest.log("Inserting data to create large chunks");
    const bigString = "X".repeat((1024 * 1024) / 4); // 1 MB
    const coll = mongos.getCollection(ns);
    let bulk = coll.initializeUnorderedBulkOp();
    for (let i = -10; i < 10; i++) {
        bulk.insert({a: i, b: i, str: bigString});
    }
    assert.commandWorked(bulk.execute());

    // Test moving large chunk with only min bound
    jsTest.log("Testing moveChunk with only min bound and large chunk");
    testMoveRangeWithBigChunk(mongos, ns, skPattern, true /* minBound */);

    // Test moving large chunk with only max bound
    jsTest.log("Testing moveChunk with only max bound and large chunk");
    testMoveRangeWithBigChunk(mongos, ns, skPattern, false /* maxBound */);
}

// Test running running moveRange on an unsplittable collection will fail
{
    const collName = "unsplittable_collection";
    const ns = kDbName + "." + collName;

    jsTest.log("Testing on unsplittable namespace");
    assert.commandWorked(mongos.getDB(kDbName).runCommand({createUnsplittableCollection: collName}));
    assert.commandFailedWithCode(
        mongos.adminCommand({moveRange: ns, min: {_id: 0}, toShard: shard0}),
        ErrorCodes.NamespaceNotSharded,
    );
}

// Test that moveRange should fail if both min and max are provided and they are not covered by a
// single chunk.
// TODO (SERVER-97588): Remove version check from tests when 9.0 becomes last LTS.
if (MongoRunner.compareBinVersions(jsTestOptions().mongosBinVersion, "8.3") >= 0) {
    assert.commandWorked(mongos.adminCommand({shardCollection: kDbName + ".collA", key: {a: 1}}));
    assert.commandWorked(mongos.adminCommand({split: kDbName + ".collA", middle: {a: 0}}));

    assert.commandFailedWithCode(
        mongos.adminCommand({
            moveRange: kDbName + ".collA",
            min: {a: MinKey()},
            max: {a: MaxKey()},
            toShard: shard1,
        }),
        11089203,
    );

    mongos.getCollection(kDbName + ".collA").drop();
}

test("nonHashedShardKey", {a: 1});

test("nonHashedCompundShardKey", {a: 1, b: 1});

test("hashedShardKey", {a: "hashed"});

st.stop();