/**
 * Tests that SBE uses the correct version of the plan cache depending on whether featureFlagSbeFull
 * is set or not.
 * @tags: [
 *   requires_profiling,
 * ]
 */
import {getLatestProfilerEntry} from "jstests/libs/profiler.js";
import {assertCacheUsage, setUpActiveCacheEntry} from "jstests/libs/query/plan_cache_utils.js";
import {checkSbeFullFeatureFlagEnabled} from "jstests/libs/query/sbe_util.js";

const conn = MongoRunner.runMongod();
const db = conn.getDB("test");
const coll = db.sbe_uses_correct_plan_cache;
const foreignCollName = "foreign";
coll.drop();

const expectedCacheVersion = checkSbeFullFeatureFlagEnabled(db) ? 2 : 1;

// assertCacheUsage() and friends require the profiler.
assert.commandWorked(db.setProfilingLevel(2));

for (let i = 0; i < 200; i++) {
    assert.commandWorked(coll.insert({a: i, b: 123}));
}

assert.commandWorked(coll.createIndex({a: 1}));
assert.commandWorked(coll.createIndex({b: 1}));

function assertCountIsOne(cursor) {
    assert.eq(1, cursor.itcount());
}

/**
 * Check that 'pipeline2' is or is not able to use a cache entry generated by 'pipeline1',
 * based on the 'shouldPipelineShareCacheEntry' flag.
 */
function assertCacheEntryIsCreatedAndUsed(
    {pipeline1, pipeline2, shouldPipelinesShareCacheEntry, cacheEntryVersion}) {
    setUpActiveCacheEntry(
        coll, pipeline1, cacheEntryVersion, "a_1" /* cachedIndexName */, assertCountIsOne);

    assertCacheUsage({
        queryColl: coll,
        pipeline: pipeline1,
        fromMultiPlanning: false,
        cacheEntryVersion,
        cacheEntryIsActive: true,
        cachedIndexName: "a_1"
    });

    coll.aggregate(pipeline2).toArray();
    assertCacheUsage({
        queryColl: coll,
        pipeline: pipeline2,
        fromMultiPlanning: !shouldPipelinesShareCacheEntry,
        cacheEntryVersion,
        cacheEntryIsActive: shouldPipelinesShareCacheEntry,
        cachedIndexName: "a_1"
    });

    coll.getPlanCache().clear();
}

// Run a simple $match/$project query and check that we use the right cache version.
assertCacheEntryIsCreatedAndUsed({
    pipeline1: [{$match: {a: 123, b: 123}}, {$project: {_id: 1, a: 1}}],
    pipeline2: [{$match: {a: 999, b: 999}}, {$project: {_id: 1, a: 1}}],
    shouldPipelinesShareCacheEntry: true,
    cacheEntryVersion: expectedCacheVersion
});

// Run two $group queries which use the same cache entry both in SBE and classic.
assertCacheEntryIsCreatedAndUsed({
    pipeline1: [{$match: {a: 123, b: 123}}, {$group: {_id: "$b", sum: {$sum: "$a"}}}],
    pipeline2: [{$match: {a: 999, b: 999}}, {$group: {_id: "$b", sum: {$sum: "$a"}}}],
    shouldPipelinesShareCacheEntry: true,
    cacheEntryVersion: expectedCacheVersion
});

// Two $group queries which can use the same classic cache entry, but not the same SBE cache entry.
// Auto-parameterization for the SBE plan cache currently only parameterizes match expressions.
assertCacheEntryIsCreatedAndUsed({
    pipeline1:
        [{$match: {a: 123, b: 123}}, {$group: {_id: "$b", sum: {$sum: {$add: [123, "$a"]}}}}],
    pipeline2:
        [{$match: {a: 123, b: 123}}, {$group: {_id: "$b", sum: {$sum: {$add: [456, "$a"]}}}}],
    shouldPipelinesShareCacheEntry: expectedCacheVersion == 1,
    cacheEntryVersion: expectedCacheVersion
});

// Check that a $match which is pushed down as part of the trailing pipeline gets parameterized,
// and can re-use a cache entry from a similar pipeline with different constants.
assertCacheEntryIsCreatedAndUsed({
    pipeline1: [
        {$match: {a: 123, b: 123}},
        {$project: {computedField: {$add: ["$a", "$b"]}}},
        {$match: {computedField: {$gt: 0}}}
    ],
    pipeline2: [
        {$match: {a: 123, b: 123}},
        {$project: {computedField: {$add: ["$a", "$b"]}}},
        {$match: {computedField: {$gt: 1}}}
    ],
    shouldPipelinesShareCacheEntry: true,
    cacheEntryVersion: expectedCacheVersion
});

// Run a pipeline with a hint and ensure that it does not get cached when the classic cache is
// being used. When the SBE cache is used, we do expect a cache entry to be written.
{
    coll.getPlanCache().clear();

    const pipe = [
        {$match: {a: 123, b: 123}},
        {$group: {_id: "$foobar"}},
    ];

    {
        const res = coll.aggregate(pipe, {hint: {a: 1}}).toArray();

        const planCacheContents = coll.getPlanCache().list();

        if (expectedCacheVersion == 1) {
            // Ensure no plan cache entry was written.
            assert.eq(planCacheContents.length, 0);
        } else if (expectedCacheVersion == 2) {
            // One entry should have been written.
            assert.eq(planCacheContents.length, 1);
        } else {
            throw 'Unknown cache version, test must be updated';
        }
    }

    // Now run the query without a hint and set up a cache entry
    setUpActiveCacheEntry(
        coll, pipe, expectedCacheVersion, "a_1" /* cachedIndexName */, assertCountIsOne);

    // Run the query with a hint.
    const res = coll.aggregate(pipe, {hint: {b: 1}}).toArray();
    const profileObj = getLatestProfilerEntry(db, {op: {$in: ["command"]}, ns: coll.getFullName()});

    // The hinted query should NOT have used the plan cache, regardless of the cache version.
    // For v:1, a hinted query should not consult the cache. For v:2, the cache key contains
    // the hint.
    assert(!profileObj.fromPlanCache);
}

MongoRunner.stopMongod(conn);
