import "jstests/libs/query/sbe_assert_error_override.js";

const coll = db[jsTestName()];
coll.drop();

const nDocs = 10;
let currentDate = new Date();
for (let i = 0; i < nDocs; i++) {
    assert.commandWorked(
        coll.insert({
            a: i,
            date: currentDate,
            array: [],
            partition: i % 2,
            partitionSeq: Math.trunc(i / 2),
        }),
    );

    const nextDate = currentDate.getDate() + 1;
    currentDate.setDate(nextDate);
}

// The list of window functions to test.
const functions = {
    sum: {$sum: "$a"},
    avg: {$avg: "$a"},
    stdDevSamp: {$stdDevSamp: "$a"},
    stdDevPop: {$stdDevPop: "$a"},
    min: {$min: "$a"},
    max: {$max: "$a"},
    count: {$count: {}},
    derivative: {$derivative: {input: "$a"}},
    derivative_date: {$derivative: {input: "$date", unit: "millisecond"}},
    integral: {$integral: {input: "$a"}},
    integral_date: {$integral: {input: "$a", unit: "millisecond"}},
    covariancePop: {$covariancePop: ["$a", "$date"]},
    covarianceSamp: {$covarianceSamp: ["$a", "$date"]},
    expMovingAvgN: {$expMovingAvg: {input: "$a", N: 3}},
    expMovingAvgAlpha: {$expMovingAvg: {input: "$a", alpha: 0.1}},
    push: {$push: "$a"},
    addToSet: {$addToSet: "$a"},
    first: {$first: "$a"},
    last: {$last: "$a"},
    shift: {$shift: {output: "$a", by: 1, default: 0}},
    documentNumber: {$documentNumber: {}},
    rank: {$rank: {}},
    denseRank: {$denseRank: {}},
};

// The list of window definitions to test.
const windows = {
    none: null,
    left_unbounded_doc: {documents: ["unbounded", "current"]},
    right_unbounded_doc: {documents: ["current", "unbounded"]},
    past_doc: {documents: [-1, "current"]},
    future_doc: {documents: ["current", 1]},
    centered_doc: {documents: [-1, 1]},
    full_unbounded_range: {range: ["unbounded", "unbounded"]},
    left_unbounded_range: {range: ["unbounded", "current"]},
    right_unbounded_range: {range: ["current", "unbounded"]},
    past_range: {range: [-2, "current"]},
    future_range: {range: ["current", 2]},
    centered_range: {range: [-2, 2]},
};

// The list of sort definitions to test.
const sortBys = {
    none: null,
    expr: {partitionSeq: {$meta: "randVal"}},
    asc: {partitionSeq: 1},
    desc: {partitionSeq: -1},
    asc_date: {date: 1},
    desc_date: {date: -1},
    multi: {partitionSeq: 1, partition: 1},
};

// The list of partition definitions to test.
const partitionBys = {
    none: null,
    field: "$partition",
    dynamic_array: "$array",
    static_array: {$const: []},
};

// Given an element from each of the lists above, construct a
// $setWindowFields stage.
function constructQuery(wf, window, sortBy, partitionBy) {
    let pathArg = Object.assign({}, wf);
    if (window != null) {
        Object.assign(pathArg, {window: window});
    }

    let setWindowFieldsArg = {};

    if (sortBy != null) {
        setWindowFieldsArg.sortBy = sortBy;
    }

    if (partitionBy != null) {
        setWindowFieldsArg.partitionBy = partitionBy;
    }

    setWindowFieldsArg.output = {x: pathArg};

    return {$setWindowFields: setWindowFieldsArg};
}

// Given an element of each of the lists above, what is the expected
// result.  The output should be 'OK' or the expected integer
// error code.
function expectedResult(wfType, windowType, sortType, partitionType) {
    // Static errors all come first.

    // Derivative and integral require an ascending sort
    // and an explicit window.
    if (wfType.startsWith("derivative")) {
        // Derivative requires a sort and an explicit window.
        if (sortType == "none" || windowType == "none") {
            return ErrorCodes.FailedToParse;
        }

        // Integral requires single column sort
        if (sortType == "multi") {
            return ErrorCodes.FailedToParse;
        }

        // '$derivative requires a non-expression sortBy'.
        if (sortType == "expr") {
            return ErrorCodes.FailedToParse;
        }
    } else if (wfType.startsWith("integral")) {
        // Integral requires a sort.
        if (sortType == "none") {
            return ErrorCodes.FailedToParse;
        }

        // Integral requires single column sort
        if (sortType == "multi") {
            return ErrorCodes.FailedToParse;
        }

        // '$integral requires a non-expression sortBy'.
        if (sortType == "expr") {
            return ErrorCodes.FailedToParse;
        }
    } else if (wfType.startsWith("expMovingAvg")) {
        // $expMovingAvg doesn't accept a window.
        if (windowType != "none") {
            return ErrorCodes.FailedToParse;
        }

        // $expMovingAvg requires a sort.
        if (sortType == "none") {
            return ErrorCodes.FailedToParse;
        }
    } else if (wfType == "documentNumber" || wfType == "rank" || wfType == "denseRank") {
        if (windowType != "none") {
            // Rank style window functions take no other arguments.
            return 5371601;
        }

        if (sortType == "none") {
            // %s must be specified with a top level sortBy expression with exactly one element.
            return 5371602;
        }

        if (sortType == "multi") {
            // %s must be specified with a top level sortBy expression with exactly one element.
            return 5371602;
        }
    } else if (wfType == "shift") {
        // $shift requires a sortBy and can't have defined a window.
        if (sortType == "none" || windowType != "none") {
            return ErrorCodes.FailedToParse;
        }
    }

    // Document windows require a sortBy.
    if (windowType.endsWith("doc") && sortType == "none") {
        // 'Document-based bounds require a sortBy'.
        return 5339901;
    }

    // Range based windows require a sort over a single field.
    if (windowType.endsWith("range")) {
        if (sortType == "none" || sortType == "multi") {
            // 'Range-based window require sortBy a single field'.
            return 5339902;
        }
        if (sortType == "expr") {
            // 'Range-based bounds require a non-expression sortBy'
            return 8947400;
        }
        if (sortType.startsWith("desc")) {
            // 'Range-based bounds require an ascending sortBy'.
            return 8947401;
        }
        if (sortType.endsWith("date") && !partitionType.endsWith("array")) {
            // 'For windows that involve date or time ranges, a unit must be provided.'
            return 5429413;
        }
    }

    if (partitionType === "static_array") {
        // When we parse $setWindowFields, we check whether partitionBy is a constant; if so we can
        // drop the partitionBy clause.  However, if the constant value is an array, we want to
        // throw an error to make it clear that partitioning by an array is not supported.
        return ErrorCodes.TypeMismatch;
    }

    // Dynamic errors all come after this point.

    if (partitionType === "dynamic_array") {
        // At runtime, we raise an error if partitionBy evaluates to an array. We chose not to
        // support partitioning by an array because $sort (which has multikey semantics)
        // doesn't partition arrays.
        return ErrorCodes.TypeMismatch;
    }

    if (wfType.startsWith("derivative")) {
        if (wfType == "derivative_date" && !sortType.endsWith("date")) {
            // "$derivative with unit expects the sortBy field to be a Date".
            return 5624900;
        }

        if (sortType.endsWith("date") && wfType != "derivative_date") {
            // "$derivative where the sortBy is a Date requires a 'unit'.
            return 5624901;
        }
    } else if (wfType.startsWith("integral")) {
        if (wfType == "integral_date" && !sortType.endsWith("date")) {
            // "$integral with unit expects the sortBy field to be a Date"
            return 5423901;
        }

        if (sortType.endsWith("date") && wfType != "integral_date") {
            // $integral where the sortBy is a Date requires a 'unit'
            return 5423902;
        }
    }

    return ErrorCodes.OK;
}

// Generate all combinations of the elements in the lists above,
// one element per list.
function* makeTests() {
    for (const [wfType, wfDefinition] of Object.entries(functions)) {
        for (const [windowType, windowDefinition] of Object.entries(windows)) {
            for (const [sortType, sortDefinition] of Object.entries(sortBys)) {
                for (const [partitionType, partitionDefinition] of Object.entries(partitionBys)) {
                    let test = {
                        query: constructQuery(wfDefinition, windowDefinition, sortDefinition, partitionDefinition),
                        expectedResult: expectedResult(wfType, windowType, sortType, partitionType),
                    };

                    yield test;
                }
            }
        }
    }
}

// Run all the combinations generated in makeTests.
for (const test of makeTests()) {
    const errorMsg = "Command was: " + tojson(test.query);
    if (test.expectedResult == ErrorCodes.OK) {
        assert.commandWorked(
            coll.runCommand({aggregate: coll.getName(), pipeline: [test.query], cursor: {}}),
            errorMsg,
        );
    } else {
        assert.commandFailedWithCode(
            coll.runCommand({aggregate: coll.getName(), pipeline: [test.query], cursor: {}}),
            test.expectedResult,
            errorMsg,
        );
    }
}