/*
 * Copyright 2022-present MongoDB, Inc.
 *
 * Licensed 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.
 */

#include "mc-fle2-rfds-private.h"

#include "mc-check-conversions-private.h"
#include "mc-fle-blob-subtype-private.h" // MC_SUBTYPE_FLE2EncryptionPlaceholder
#include "mongocrypt-private.h"          // CLIENT_ERR

#include "mlib/thread.h" // mlib_once_flag
#include <math.h>        // INFINITY

static mc_FLE2RangeOperator_t get_operator_type(const char *key) {
    BSON_ASSERT_PARAM(key);

    if (0 == strcmp(key, "$gt")) {
        return FLE2RangeOperator_kGt;
    } else if (0 == strcmp(key, "$gte")) {
        return FLE2RangeOperator_kGte;
    } else if (0 == strcmp(key, "$lt")) {
        return FLE2RangeOperator_kLt;
    } else if (0 == strcmp(key, "$lte")) {
        return FLE2RangeOperator_kLte;
    } else {
        return FLE2RangeOperator_kNone;
    }
}

static const char *mc_FLE2RangeOperator_to_string(mc_FLE2RangeOperator_t op) {
    switch (op) {
    case FLE2RangeOperator_kGt: return "$gt";
    case FLE2RangeOperator_kGte: return "$gte";
    case FLE2RangeOperator_kLt: return "$lt";
    case FLE2RangeOperator_kLte: return "$lte";
    case FLE2RangeOperator_kNone: return "none";
    default: return "Unknown";
    }
}

static bool is_supported_operator(const char *key) {
    BSON_ASSERT_PARAM(key);
    return get_operator_type(key) != FLE2RangeOperator_kNone;
}

// Suffix an error message with the JSON string form of bson.
#define ERR_WITH_BSON(bson, fmt, ...)                                                                                  \
    if (1) {                                                                                                           \
        char *_str = bson_as_canonical_extended_json((bson), NULL);                                                    \
        CLIENT_ERR(fmt ": %s", __VA_ARGS__, _str);                                                                     \
        bson_free(_str);                                                                                               \
    } else                                                                                                             \
        (void)0

// Parses a document like {$and: []} and outputs an iterator to the array.
static bool parse_and(const bson_t *in, bson_iter_t *out, mongocrypt_status_t *status) {
    BSON_ASSERT_PARAM(in);
    BSON_ASSERT_PARAM(out);
    BSON_ASSERT(status || true);

    bson_iter_t and = {0};
    if (!bson_iter_init(&and, in) || !bson_iter_next(&and) || 0 != strcmp(bson_iter_key(&and), "$and")) {
        ERR_WITH_BSON(in, "%s", "error unable to find '$and'");
        return false;
    }

    if (!BSON_ITER_HOLDS_ARRAY(&and)) {
        ERR_WITH_BSON(in, "%s", "expected '$and' to be array");
        return false;
    }

    *out = and;

    if (bson_iter_next(&and)) {
        ERR_WITH_BSON(in, "unexpected extra key '%s' after '$and'", bson_iter_key(&and));
        return false;
    }
    return true;
}

typedef struct {
    const char *op_type_str;
    const char *field;
    bson_iter_t value;
    mc_FLE2RangeOperator_t op_type;
} operator_value_t;

// Parses a document like {$gt: ["$age", 5]}.
static bool
parse_aggregate_expression(const bson_t *orig, bson_iter_t *in, operator_value_t *out, mongocrypt_status_t *status) {
    BSON_ASSERT_PARAM(orig);
    BSON_ASSERT_PARAM(in);
    BSON_ASSERT_PARAM(out);
    BSON_ASSERT(status || true);

    bson_iter_t array = {0}, value;
    const char *op_type_str = bson_iter_key(in);
    bool ok = false;
    const char *field;

    if (!BSON_ITER_HOLDS_ARRAY(in)) {
        ERR_WITH_BSON(orig, "%s", "expected argument to be array");
        goto fail;
    }

    if (!bson_iter_recurse(in, &array)) {
        ERR_WITH_BSON(orig, "%s", "failed to recurse into array");
        goto fail;
    }

    // Expect exactly 2 elements, like ["$age", 5]. The first element is the
    // field path. The second element is the value.
    if (!bson_iter_next(&array)) {
        ERR_WITH_BSON(orig, "expected 2 elements in operand %s, got 0", op_type_str);
        goto fail;
    }
    if (!BSON_ITER_HOLDS_UTF8(&array)) {
        ERR_WITH_BSON(orig, "expected UTF-8 as first element in %s", op_type_str);
        goto fail;
    }
    field = bson_iter_utf8(&array, NULL);

    if (!bson_iter_next(&array)) {
        ERR_WITH_BSON(orig, "expected 2 elements in operand %s, got 1", op_type_str);
        goto fail;
    }
    value = array;
    if (bson_iter_next(&array)) {
        ERR_WITH_BSON(orig, "expected 2 elements in operand %s, got > 2", op_type_str);
        goto fail;
    }

    *out = (operator_value_t){
        .field = field,
        .op_type_str = op_type_str,
        .op_type = get_operator_type(op_type_str),
        .value = value,
    };

    ok = true;
fail:
    return ok;
}

// Parses a document like {age: {$gt: 5}}.
static bool
parse_match_expression(const bson_t *orig, bson_iter_t *in, operator_value_t *out, mongocrypt_status_t *status) {
    BSON_ASSERT_PARAM(orig);
    BSON_ASSERT_PARAM(in);
    BSON_ASSERT_PARAM(out);
    BSON_ASSERT(status || true);

    bson_iter_t document = {0}, value;
    const char *op_type_str;
    bool ok = false;
    const char *field = bson_iter_key(in);

    if (!BSON_ITER_HOLDS_DOCUMENT(in)) {
        ERR_WITH_BSON(orig, "%s", "expected argument to be document");
        goto fail;
    }

    if (!bson_iter_recurse(in, &document)) {
        ERR_WITH_BSON(orig, "%s", "failed to recurse into document");
        goto fail;
    }

    // Expect exactly 1 elements, like {$gt: 5}.
    if (!bson_iter_next(&document)) {
        ERR_WITH_BSON(orig, "expected 1 elements in operand %s, got 0", field);
        goto fail;
    }
    op_type_str = bson_iter_key(&document);
    if (!is_supported_operator(op_type_str)) {
        ERR_WITH_BSON(orig, "unsupported operator: %s", op_type_str);
        goto fail;
    }
    value = document;

    if (bson_iter_next(&document)) {
        ERR_WITH_BSON(orig, "expected 1 elements in operand %s, got > 1", field);
        goto fail;
    }

    *out = (operator_value_t){
        .field = field,
        .op_type_str = op_type_str,
        .op_type = get_operator_type(op_type_str),
        .value = value,
    };

    ok = true;
fail:
    return ok;
}

bool mc_FLE2RangeFindDriverSpec_parse(mc_FLE2RangeFindDriverSpec_t *spec,
                                      const bson_t *in,
                                      mongocrypt_status_t *status) {
    BSON_ASSERT_PARAM(in);
    BSON_ASSERT(status || true);

    *spec = (mc_FLE2RangeFindDriverSpec_t){0};
    // `in` may be an Aggregate Expression with this form:
    // {$and: [{$gt: ["$age", 5]}, {$lt:["$age", 50]}]}
    // Or `in` may be a Match Expression with this form:
    // {$and: [{age: {$gt: 5}}, {age: {$lt: 50}} ]}
    bson_iter_t and = {0}, array = {0};
    bool ok = false;

    if (!parse_and(in, &and, status)) {
        goto fail;
    }

    // Iterate over array elements.
    if (!bson_iter_recurse(&and, &array)) {
        ERR_WITH_BSON(in, "%s", "failed to recurse into '$and'");
        goto fail;
    }

    enum { UNKNOWN, MATCH_EXPRESSION, AGGREGATE_EXPRESSION } arg_type = UNKNOWN;

    while (bson_iter_next(&array)) {
        bson_iter_t doc;

        if (!BSON_ITER_HOLDS_DOCUMENT(&array)) {
            ERR_WITH_BSON(in, "%s", "expected document in '$and' array");
            goto fail;
        }

        if (!bson_iter_recurse(&array, &doc)) {
            ERR_WITH_BSON(in, "%s", "failed to recurse into '$and' element");
            goto fail;
        }

        if (!bson_iter_next(&doc)) {
            ERR_WITH_BSON(in, "%s", "unexpected empty '$and' array document");
            goto fail;
        }

        if (arg_type == UNKNOWN) {
            // Attempt to determine argument type by inspecting first key.
            if (is_supported_operator(bson_iter_key(&doc))) {
                // Assume the document is part of an Aggregate Expression, like:
                // {$gt: ["$age", 5]}
                arg_type = AGGREGATE_EXPRESSION;
                spec->isAggregateExpression = true;
            } else {
                // Assume the document is part of a Match Expression, like:
                // {age: {$gt: 5}}
                arg_type = MATCH_EXPRESSION;
            }
        }

        operator_value_t op = {0};
        switch (arg_type) {
        case AGGREGATE_EXPRESSION:
            if (!parse_aggregate_expression(in, &doc, &op, status)) {
                goto fail;
            }
            break;
        case MATCH_EXPRESSION:
            if (!parse_match_expression(in, &doc, &op, status)) {
                goto fail;
            }
            break;
        case UNKNOWN:
        default: ERR_WITH_BSON(in, "%s", "unexpected unknown expression type"); goto fail;
        }

        switch (op.op_type) {
        case FLE2RangeOperator_kGt:
        case FLE2RangeOperator_kGte:
            if (spec->lower.set) {
                ERR_WITH_BSON(in, "unexpected duplicate bound %s", op.op_type_str);
                goto fail;
            }
            spec->lower.set = true;
            spec->lower.value = op.value;
            spec->lower.included = op.op_type == FLE2RangeOperator_kGte;
            break;
        case FLE2RangeOperator_kLt:
        case FLE2RangeOperator_kLte:
            if (spec->upper.set) {
                ERR_WITH_BSON(in, "unexpected duplicate bound %s", op.op_type_str);
                goto fail;
            }
            spec->upper.set = true;
            spec->upper.value = op.value;
            spec->upper.included = op.op_type == FLE2RangeOperator_kLte;
            break;
        case FLE2RangeOperator_kNone:
        default:
            ERR_WITH_BSON(in, "unsupported operator type %s", op.op_type_str);
            goto fail;
            break;
        }

        if (spec->field && 0 != strcmp(spec->field, op.field)) {
            ERR_WITH_BSON(in,
                          "unexpected field mismatch. Expected all fields to be "
                          "%s, but got %s",
                          spec->field,
                          op.field);
            goto fail;
        }

        spec->nOps++;

        switch (spec->nOps) {
        case 1: spec->firstOp = op.op_type; break;
        case 2: spec->secondOp = op.op_type; break;
        default: ERR_WITH_BSON(in, "expected 1 or 2 operators, got > 2: %s", op.op_type_str); goto fail;
        }

        spec->field = op.field;
    }

    ok = true;
fail:
    return ok;
}

// mc_makeRangeFindPlaceholder creates a placeholder to be consumed by
// libmongocrypt for encryption.
bool mc_makeRangeFindPlaceholder(mc_makeRangeFindPlaceholder_args_t *args,
                                 _mongocrypt_buffer_t *out,
                                 mongocrypt_status_t *status) {
    BSON_ASSERT_PARAM(args);
    BSON_ASSERT_PARAM(out);
    BSON_ASSERT(status || true);

    bool ok = false;
    bson_t *edgesInfo = bson_new();
    bson_t *v = bson_new();
    bson_t *p = bson_new();
    _mongocrypt_buffer_init(out);

#define TRY(stmt)                                                                                                      \
    if (!(stmt)) {                                                                                                     \
        CLIENT_ERR("error appending BSON for placeholder");                                                            \
        goto fail;                                                                                                     \
    }

    // create edgesInfo.

    if (!args->isStub) {
        TRY(bson_append_iter(edgesInfo, "lowerBound", -1, &args->lowerBound));
        TRY(BSON_APPEND_BOOL(edgesInfo, "lbIncluded", args->lbIncluded));
        TRY(bson_append_iter(edgesInfo, "upperBound", -1, &args->upperBound));
        TRY(BSON_APPEND_BOOL(edgesInfo, "ubIncluded", args->ubIncluded));
        TRY(bson_append_iter(edgesInfo, "indexMin", -1, &args->indexMin));
        TRY(bson_append_iter(edgesInfo, "indexMax", -1, &args->indexMax));
        if (args->precision.set) {
            TRY(BSON_APPEND_INT32(edgesInfo, "precision", args->precision.value));
        }
        if (args->trimFactor.set) {
            TRY(BSON_APPEND_INT32(edgesInfo, "trimFactor", args->trimFactor.value));
        }
        TRY(BSON_APPEND_DOCUMENT(v, "edgesInfo", edgesInfo));
    }

    // create v.
    TRY(BSON_APPEND_INT32(v, "payloadId", args->payloadId));
    TRY(BSON_APPEND_INT32(v, "firstOperator", (int32_t)args->firstOp));
    if (args->secondOp) {
        TRY(BSON_APPEND_INT32(v, "secondOperator", (int32_t)args->secondOp));
    }

    // create placeholder.
    TRY(BSON_APPEND_INT32(p, "t", MONGOCRYPT_FLE2_PLACEHOLDER_TYPE_FIND));
    TRY(BSON_APPEND_INT32(p, "a", MONGOCRYPT_FLE2_ALGORITHM_RANGE));
    TRY(_mongocrypt_buffer_append(args->index_key_id, p, "ki", 2));
    TRY(_mongocrypt_buffer_append(args->user_key_id, p, "ku", 2));
    TRY(BSON_APPEND_DOCUMENT(p, "v", v));
    TRY(BSON_APPEND_INT64(p, "cm", args->maxContentionFactor));
    TRY(BSON_APPEND_INT64(p, "s", args->sparsity));
#undef TRY

    BSON_ASSERT(p->len < UINT32_MAX);
    _mongocrypt_buffer_resize(out, p->len + 1u);
    out->subtype = BSON_SUBTYPE_ENCRYPTED;
    out->data[0] = MC_SUBTYPE_FLE2EncryptionPlaceholder;
    memcpy(out->data + 1, bson_get_data(p), p->len);

    ok = true;
fail:
    bson_destroy(p);
    bson_destroy(v);
    bson_destroy(edgesInfo);
    return ok;
}

bool mc_FLE2RangeFindDriverSpec_to_placeholders(mc_FLE2RangeFindDriverSpec_t *spec,
                                                const mc_RangeOpts_t *range_opts,
                                                int64_t maxContentionFactor,
                                                const _mongocrypt_buffer_t *user_key_id,
                                                const _mongocrypt_buffer_t *index_key_id,
                                                int32_t payloadId,
                                                bson_t *out,
                                                mongocrypt_status_t *status) {
    BSON_ASSERT_PARAM(spec);
    BSON_ASSERT_PARAM(range_opts);
    BSON_ASSERT_PARAM(user_key_id);
    BSON_ASSERT_PARAM(index_key_id);
    BSON_ASSERT_PARAM(out);
    BSON_ASSERT(status || true);

    _mongocrypt_buffer_t p1 = {0}, p2 = {0};
    bson_t infDoc = BSON_INITIALIZER;
    bson_iter_t negInf, posInf;
    bson_t minMaxDoc = BSON_INITIALIZER;
    bson_iter_t indexMin, indexMax;
    bool ok = false;

    BCON_APPEND(&infDoc, "p", BCON_DOUBLE(INFINITY), "n", BCON_DOUBLE(-INFINITY));

#define TRY(stmt)                                                                                                      \
    if (!(stmt)) {                                                                                                     \
        CLIENT_ERR("error transforming BSON for FLE2RangeFindDriverSpec: %s", #stmt);                                  \
        goto fail;                                                                                                     \
    }

    TRY(bson_iter_init_find(&posInf, &infDoc, "p"));
    TRY(bson_iter_init_find(&negInf, &infDoc, "n"));

    // Apply default index min/max values.
    {
        bson_type_t index_type;
        if (spec->lower.set) {
            index_type = bson_iter_type(&spec->lower.value);
        } else if (spec->upper.set) {
            index_type = bson_iter_type(&spec->upper.value);
        } else {
            CLIENT_ERR("expected lower or upper bound to be set");
            goto fail;
        }
        if (!mc_RangeOpts_appendMin(range_opts, index_type, "indexMin", &minMaxDoc, status)) {
            goto fail;
        }
        if (!mc_RangeOpts_appendMax(range_opts, index_type, "indexMax", &minMaxDoc, status)) {
            goto fail;
        }

        TRY(bson_iter_init_find(&indexMin, &minMaxDoc, "indexMin"));
        TRY(bson_iter_init_find(&indexMax, &minMaxDoc, "indexMax"));
    }

    mc_makeRangeFindPlaceholder_args_t args = {.isStub = false,
                                               .user_key_id = user_key_id,
                                               .index_key_id = index_key_id,
                                               .lowerBound = spec->lower.set ? spec->lower.value : negInf,
                                               .lbIncluded = spec->lower.set ? spec->lower.included : true,
                                               .upperBound = spec->upper.set ? spec->upper.value : posInf,
                                               .ubIncluded = spec->upper.set ? spec->upper.included : true,
                                               .payloadId = payloadId,
                                               .firstOp = spec->firstOp,
                                               .secondOp = spec->secondOp,
                                               .indexMin = indexMin,
                                               .indexMax = indexMax,
                                               .precision = range_opts->precision,
                                               .maxContentionFactor = maxContentionFactor,
                                               .sparsity = range_opts->sparsity,
                                               .trimFactor = range_opts->trimFactor};

    // First operator is the non-stub.
    if (!mc_makeRangeFindPlaceholder(&args, &p1, status)) {
        goto fail;
    }

    // Second operator (if required) is a stub.
    if (spec->nOps == 2) {
        mc_makeRangeFindPlaceholder_args_t args = {.isStub = true,
                                                   .user_key_id = user_key_id,
                                                   .index_key_id = index_key_id,
                                                   .payloadId = payloadId,
                                                   .firstOp = spec->firstOp,
                                                   .secondOp = spec->secondOp,
                                                   .maxContentionFactor = maxContentionFactor,
                                                   .sparsity = range_opts->sparsity};

        // First operator is the non-stub.
        if (!mc_makeRangeFindPlaceholder(&args, &p2, status)) {
            goto fail;
        }
    }

    if (spec->isAggregateExpression) {
        /*
        Create an Aggregate Expression document like:
        {
        "$and" : [
              {"$gt" : [ "$age", "<placeholder1>" ]},
              {"$lt" : [ "$age", "<placeholder2>" ]}
           ]
        }
        */
        bson_t and;
        TRY(BSON_APPEND_ARRAY_BEGIN(out, "$and", &and));
        bson_t elem;
        TRY(BSON_APPEND_DOCUMENT_BEGIN(&and, "0", &elem));
        bson_t operator;
        TRY(BSON_APPEND_ARRAY_BEGIN(&elem, mc_FLE2RangeOperator_to_string(spec->firstOp), &operator));
        TRY(BSON_APPEND_UTF8(&operator, "0", spec->field));
        TRY(_mongocrypt_buffer_append(&p1, &operator, "1", 1));
        TRY(bson_append_array_end(&elem, &operator));
        TRY(bson_append_document_end(&and, &elem));

        if (spec->nOps == 2) {
            TRY(BSON_APPEND_DOCUMENT_BEGIN(&and, "1", &elem));
            TRY(BSON_APPEND_ARRAY_BEGIN(&elem, mc_FLE2RangeOperator_to_string(spec->secondOp), &operator));
            TRY(BSON_APPEND_UTF8(&operator, "0", spec->field));
            TRY(_mongocrypt_buffer_append(&p2, &operator, "1", 1));
            TRY(bson_append_array_end(&elem, &operator));
            TRY(bson_append_document_end(&and, &elem));
        }

        TRY(bson_append_array_end(out, &and));
    } else {
        /*
        Create a Match Expression document like:
        {
          "$and" : [
              {"age" : { "$gt": "<placeholder1>" }},
              {"age" : { "$lt": "<placeholder2>" }}
           ]
        }
        */
        bson_t and;
        TRY(BSON_APPEND_ARRAY_BEGIN(out, "$and", &and));
        bson_t elem;
        TRY(BSON_APPEND_DOCUMENT_BEGIN(&and, "0", &elem));
        bson_t operator;
        TRY(BSON_APPEND_DOCUMENT_BEGIN(&elem, spec->field, &operator));
        const char *op_str = mc_FLE2RangeOperator_to_string(spec->firstOp);
        TRY(_mongocrypt_buffer_append(&p1, &operator, op_str, -1));
        TRY(bson_append_document_end(&elem, &operator));
        TRY(bson_append_document_end(&and, &elem));

        if (spec->nOps == 2) {
            TRY(BSON_APPEND_DOCUMENT_BEGIN(&and, "1", &elem));
            TRY(BSON_APPEND_DOCUMENT_BEGIN(&elem, spec->field, &operator));
            const char *op_str = mc_FLE2RangeOperator_to_string(spec->secondOp);
            TRY(_mongocrypt_buffer_append(&p2, &operator, op_str, -1));
            TRY(bson_append_document_end(&elem, &operator));
            TRY(bson_append_document_end(&and, &elem));
        }

        TRY(bson_append_array_end(out, &and));
    }

#undef TRY
    ok = true;
fail:
    _mongocrypt_buffer_cleanup(&p2);
    _mongocrypt_buffer_cleanup(&p1);
    bson_destroy(&minMaxDoc);
    bson_destroy(&infDoc);
    return ok;
}

static mlib_once_flag payloadId_init_flag = MLIB_ONCE_INITIALIZER;
static mongocrypt_mutex_t payloadId_mutex;
static int32_t payloadId = 0;

static void payloadId_init_mutex(void) {
    _mongocrypt_mutex_init(&payloadId_mutex);
}

void mc_reset_payloadId_for_testing(void) {
    mlib_call_once(&payloadId_init_flag, payloadId_init_mutex);
    MONGOCRYPT_WITH_MUTEX(payloadId_mutex) {
        payloadId = 0;
    }
}

// mc_getNextPayloadId is thread safe.
int32_t mc_getNextPayloadId(void) {
    mlib_call_once(&payloadId_init_flag, payloadId_init_mutex);
    int32_t ret;
    MONGOCRYPT_WITH_MUTEX(payloadId_mutex) {
        ret = payloadId;
        payloadId = payloadId < INT32_MAX ? payloadId + 1 : 0;
    }
    return ret;
}
