ShaiBot/utils/cookieEconomy.js

const {
    MessageFlags
} = require("discord.js");

const {
    mClient
} = require("../index");


/*
 * ============================================================
 * DATABASE CONFIG
 * ============================================================
 *
 * Set M_DB in .env to the name of your Mongo database.
 *
 * Example:
 *
 * M_DB=ShaiBot
 */

const DATABASE_NAME =
    process.env.M_DB;


/*
 * Cookie balances:
 *
 * {
 *     userId: "123456789",
 *     cookies: 2500
 * }
 */

const COOKIE_COLLECTION =
    "items_cookies";


/*
 * Cooldowns:
 *
 * {
 *     userId: "123456789",
 *     guildId: "987654321",
 *
 *     collect: 1787000000,
 *     steal: 1787000000,
 *     heist: 1787000000
 * }
 *
 * Values are Unix timestamps representing when
 * the action becomes available again.
 *
 * No createdAt / updatedAt fields.
 */

const COOLDOWN_COLLECTION =
    "cooldown_cookies";


/*
 * ============================================================
 * COLLECTION ACCESS
 * ============================================================
 */

function getCollections() {
    if (!DATABASE_NAME) {
        throw new Error(
            "M_DB is not configured in .env"
        );
    }


    const db =
        mClient.db(
            DATABASE_NAME
        );


    return {
        db,

        cookies:
            db.collection(
                COOKIE_COLLECTION
            ),

        cooldowns:
            db.collection(
                COOLDOWN_COLLECTION
            )
    };
}


/*
 * ============================================================
 * UNIX TIME
 * ============================================================
 */

function getUnixTime() {
    return Math.floor(
        Date.now() /
        1000
    );
}


/*
 * ============================================================
 * COOKIE BALANCE
 * ============================================================
 */

async function getCookieBalance(
    userId
) {
    const {
        cookies
    } =
        getCollections();


    const document =
        await cookies.findOne({
            userId
        });


    return Math.max(
        0,
        document?.cookies ??
        0
    );
}


/*
 * ============================================================
 * CHANGE COOKIE BALANCE
 * ============================================================
 *
 * Positive:
 *
 * changeCookies(userId, 50)
 *
 * Negative:
 *
 * changeCookies(userId, -50)
 *
 * Negative changes are protected against taking a
 * balance below zero.
 */

async function changeCookies(
    userId,
    amount
) {
    if (
        !Number.isSafeInteger(
            amount
        )
    ) {
        throw new Error(
            "INVALID_COOKIE_AMOUNT"
        );
    }


    if (
        amount ===
        0
    ) {
        return getCookieBalance(
            userId
        );
    }


    const {
        cookies
    } =
        getCollections();


    /*
     * Adding Cookies.
     */

    if (
        amount >
        0
    ) {
        await cookies.updateOne(
            {
                userId
            },
            {
                $inc: {
                    cookies:
                        amount
                },

                $setOnInsert: {
                    userId
                }
            },
            {
                upsert:
                    true
            }
        );


        return getCookieBalance(
            userId
        );
    }


    /*
     * Removing Cookies.
     *
     * Atomic balance check prevents the account
     * from becoming negative.
     */

    const removeAmount =
        Math.abs(
            amount
        );


    const result =
        await cookies.updateOne(
            {
                userId,

                cookies: {
                    $gte:
                        removeAmount
                }
            },
            {
                $inc: {
                    cookies:
                        -removeAmount
                }
            }
        );


    if (
        result.modifiedCount !==
        1
    ) {
        throw new Error(
            "INSUFFICIENT_COOKIES"
        );
    }


    return getCookieBalance(
        userId
    );
}


/*
 * ============================================================
 * IDEMPOTENT COOKIE CHANGE
 * ============================================================
 *
 * Applies a Cookie mutation at most once for a user-scoped
 * operationId. This is useful for persistent game buy-ins and
 * cash-outs that may be retried after a crash or restart.
 *
 * The Cookie balance mutation and operation marker are written
 * atomically to the same user document.
 */

async function changeCookiesOnce(
    userId,
    amount,
    operationId
) {
    if (
        !Number.isSafeInteger(
            amount
        )
    ) {
        throw new Error(
            "INVALID_COOKIE_AMOUNT"
        );
    }

    if (
        typeof operationId !==
            "string"
        ||
        !/^[A-Za-z0-9:_-]{1,160}$/.test(
            operationId
        )
    ) {
        throw new Error(
            "INVALID_COOKIE_OPERATION_ID"
        );
    }

    if (
        amount ===
        0
    ) {
        return {
            applied:
                false,
            balance:
                await getCookieBalance(
                    userId
                )
        };
    }

    const {
        cookies
    } =
        getCollections();

    const filter = {
        userId,
        appliedCookieOperations: {
            $ne:
                operationId
        }
    };

    if (
        amount <
        0
    ) {
        filter.cookies = {
            $gte:
                Math.abs(
                    amount
                )
        };
    }

    const update = {
        $inc: {
            cookies:
                amount
        },
        $addToSet: {
            appliedCookieOperations:
                operationId
        }
    };

    if (
        amount >
        0
    ) {
        update.$setOnInsert = {
            userId
        };
    }

    try {
        const result =
            await cookies.updateOne(
                filter,
                update,
                {
                    upsert:
                        amount > 0
                }
            );

        if (
            result.modifiedCount ===
                1
            ||
            result.upsertedCount ===
                1
        ) {
            return {
                applied:
                    true,
                balance:
                    await getCookieBalance(
                        userId
                    )
            };
        }
    }

    catch (error) {
        /*
         * A concurrent positive retry can race an upsert.
         * The unique userId index may reject the losing insert.
         * Re-read below and treat an already-recorded operation
         * as a successful idempotent retry.
         */
        if (
            error?.code !==
            11000
        ) {
            throw error;
        }
    }

    const document =
        await cookies.findOne(
            {
                userId
            },
            {
                projection: {
                    cookies:
                        1,
                    appliedCookieOperations:
                        1
                }
            }
        );

    const operations =
        Array.isArray(
            document
                ?.appliedCookieOperations
        )
            ? document
                .appliedCookieOperations
            : [];

    if (
        operations.includes(
            operationId
        )
    ) {
        return {
            applied:
                false,
            balance:
                Math.max(
                    0,
                    Number(
                        document?.cookies ??
                        0
                    )
                )
        };
    }

    if (
        amount <
        0
    ) {
        throw new Error(
            "INSUFFICIENT_COOKIES"
        );
    }

    throw new Error(
        "COOKIE_OPERATION_FAILED"
    );
}


async function clearCookieOperation(
    userId,
    operationId
) {
    if (
        typeof operationId !==
            "string"
        ||
        !/^[A-Za-z0-9:_-]{1,160}$/.test(
            operationId
        )
    ) {
        throw new Error(
            "INVALID_COOKIE_OPERATION_ID"
        );
    }

    const {
        cookies
    } =
        getCollections();

    const result =
        await cookies.updateOne(
            {
                userId
            },
            {
                $pull: {
                    appliedCookieOperations:
                        operationId
                }
            }
        );

    /*
     * Keep the Cookie document tidy once the final temporary
     * idempotency marker has been cleared.
     */

    await cookies.updateOne(
        {
            userId,
            appliedCookieOperations: {
                $size: 0
            }
        },
        {
            $unset: {
                appliedCookieOperations:
                    ""
            }
        }
    );

    return result.modifiedCount ===
        1;
}


async function hasCookieOperation(
    userId,
    operationId
) {
    if (
        typeof operationId !==
            "string"
        ||
        !/^[A-Za-z0-9:_-]{1,160}$/.test(
            operationId
        )
    ) {
        throw new Error(
            "INVALID_COOKIE_OPERATION_ID"
        );
    }

    const {
        cookies
    } =
        getCollections();

    const document =
        await cookies.findOne(
            {
                userId,
                appliedCookieOperations:
                    operationId
            },
            {
                projection: {
                    _id:
                        1
                }
            }
        );

    return Boolean(
        document
    );
}


/*
 * ============================================================
 * COOKIE TRANSFER
 * ============================================================
 *
 * This intentionally DOES NOT use Mongo transactions.
 *
 * Your Mongo deployment does not support transactions /
 * retryable transactional writes.
 *
 * The debit is atomic.
 *
 * If crediting the destination fails after the debit,
 * we attempt to compensate by returning the Cookies
 * to the source account.
 */

async function transferCookies(
    fromUserId,
    toUserId,
    amount
) {
    if (
        fromUserId ===
        toUserId
    ) {
        throw new Error(
            "SAME_ACCOUNT"
        );
    }


    if (
        !Number.isSafeInteger(
            amount
        )
        ||
        amount <=
        0
    ) {
        throw new Error(
            "INVALID_COOKIE_AMOUNT"
        );
    }


    const {
        cookies
    } =
        getCollections();


    /*
     * ========================================================
     * DEBIT SOURCE
     * ========================================================
     *
     * This is the important atomic operation.
     *
     * Two simultaneous attempts cannot both spend the
     * same Cookies unless the balance actually covers both.
     */

    const debit =
        await cookies.updateOne(
            {
                userId:
                    fromUserId,

                cookies: {
                    $gte:
                        amount
                }
            },
            {
                $inc: {
                    cookies:
                        -amount
                }
            }
        );


    if (
        debit.modifiedCount !==
        1
    ) {
        throw new Error(
            "INSUFFICIENT_COOKIES"
        );
    }


    /*
     * ========================================================
     * CREDIT DESTINATION
     * ========================================================
     */

    try {
        await cookies.updateOne(
            {
                userId:
                    toUserId
            },
            {
                $inc: {
                    cookies:
                        amount
                },

                $setOnInsert: {
                    userId:
                        toUserId
                }
            },
            {
                upsert:
                    true
            }
        );

    } catch (error) {
        /*
         * Best-effort compensation.
         *
         * If destination credit fails, return the
         * Cookies we already removed.
         */

        try {
            await cookies.updateOne(
                {
                    userId:
                        fromUserId
                },
                {
                    $inc: {
                        cookies:
                            amount
                    },

                    $setOnInsert: {
                        userId:
                            fromUserId
                    }
                },
                {
                    upsert:
                        true
                }
            );

        } catch (refundError) {
            console.error(
                "CRITICAL: Cookie transfer compensation failed:",
                {
                    fromUserId,
                    toUserId,
                    amount,
                    refundError
                }
            );
        }


        throw error;
    }


    return {
        amount,

        from:
            fromUserId,

        to:
            toUserId
    };
}


/*
 * ============================================================
 * GET COOLDOWN
 * ============================================================
 *
 * Returns the Unix timestamp when the cooldown becomes
 * available again.
 *
 * Returns 0 when no cooldown exists.
 */

async function getCooldown(
    userId,
    guildId,
    cooldownName
) {
    validateCooldownName(
        cooldownName
    );


    const {
        cooldowns
    } =
        getCollections();


    const document =
        await cooldowns.findOne(
            {
                userId,
                guildId
            },
            {
                projection: {
                    [cooldownName]:
                        1
                }
            }
        );


    return Number(
        document?.[
            cooldownName
        ]
        ??
        0
    );
}


/*
 * ============================================================
 * SET COOLDOWN
 * ============================================================
 *
 * durationSeconds is added to the current Unix timestamp.
 *
 * Returns the timestamp when the action becomes available.
 */

async function setCooldown(
    userId,
    guildId,
    cooldownName,
    durationSeconds
) {
    validateCooldownName(
        cooldownName
    );


    validateCooldownDuration(
        durationSeconds
    );


    const {
        cooldowns
    } =
        getCollections();


    const next =
        getUnixTime() +
        durationSeconds;


    await cooldowns.updateOne(
        {
            userId,
            guildId
        },
        {
            $set: {
                [cooldownName]:
                    next
            },

            $setOnInsert: {
                userId,
                guildId
            }
        },
        {
            upsert:
                true
        }
    );


    return next;
}


/*
 * ============================================================
 * CLAIM COOLDOWN
 * ============================================================
 *
 * Used for actions where rapid double clicks must not be
 * allowed to resolve twice.
 *
 * Vault Heist uses:
 *
 * claimCooldown(
 *     userId,
 *     guildId,
 *     "heist",
 *     HEIST_COOLDOWN
 * )
 *
 * Sneak / Hack / Smash therefore all share ONE cooldown.
 *
 * Returns:
 *
 * {
 *     claimed: true,
 *     next: 1787000000
 * }
 *
 * or:
 *
 * {
 *     claimed: false,
 *     next: 1787000000
 * }
 */

async function claimCooldown(
    userId,
    guildId,
    cooldownName,
    durationSeconds
) {
    validateCooldownName(
        cooldownName
    );


    validateCooldownDuration(
        durationSeconds
    );


    const {
        cooldowns
    } =
        getCollections();


    const now =
        getUnixTime();


    const next =
        now +
        durationSeconds;


    /*
     * First attempt:
     *
     * Update an existing cooldown document only if
     * this particular cooldown is absent or expired.
     *
     * This is atomic.
     */

    const existing =
        await cooldowns.updateOne(
            {
                userId,
                guildId,

                $or: [
                    {
                        [cooldownName]: {
                            $exists:
                                false
                        }
                    },

                    {
                        [cooldownName]: {
                            $lte:
                                now
                        }
                    }
                ]
            },
            {
                $set: {
                    [cooldownName]:
                        next
                }
            }
        );


    if (
        existing.modifiedCount ===
        1
    ) {
        return {
            claimed:
                true,

            next
        };
    }


    /*
     * The document might not exist at all yet.
     *
     * Try inserting one.
     *
     * If another interaction created it at the exact
     * same moment, the unique index will reject one
     * side and we'll simply read the winning cooldown.
     */

    try {
        await cooldowns.insertOne({
            userId,
            guildId,

            [cooldownName]:
                next
        });


        return {
            claimed:
                true,

            next
        };

    } catch (error) {
        /*
         * 11000 = duplicate key.
         *
         * That is expected when two interactions race
         * to create the same cooldown record.
         */

        if (
            error?.code !==
            11000
        ) {
            throw error;
        }
    }


    /*
     * Someone else already owns the active cooldown.
     */

    const current =
        await getCooldown(
            userId,
            guildId,
            cooldownName
        );


    /*
     * Rare boundary case:
     *
     * the cooldown expired between our previous write
     * attempt and this read.
     *
     * Try one more atomic claim.
     */

    if (
        current <=
        getUnixTime()
    ) {
        const retry =
            await cooldowns.updateOne(
                {
                    userId,
                    guildId,

                    [cooldownName]: {
                        $lte:
                            getUnixTime()
                    }
                },
                {
                    $set: {
                        [cooldownName]:
                            next
                    }
                }
            );


        if (
            retry.modifiedCount ===
            1
        ) {
            return {
                claimed:
                    true,

                next
            };
        }
    }


    return {
        claimed:
            false,

        next:
            current
    };
}


/*
 * ============================================================
 * CLEAR COOLDOWN
 * ============================================================
 *
 * Useful during testing/admin work.
 */

async function clearCooldown(
    userId,
    guildId,
    cooldownName
) {
    validateCooldownName(
        cooldownName
    );


    const {
        cooldowns
    } =
        getCollections();


    await cooldowns.updateOne(
        {
            userId,
            guildId
        },
        {
            $unset: {
                [cooldownName]:
                    ""
            }
        }
    );
}


/*
 * ============================================================
 * TEMPORARY EPHEMERAL REPLY
 * ============================================================
 *
 * Usage:
 *
 * await temporaryEphemeralReply(
 *     interaction,
 *     {
 *         content: "Nope."
 *     },
 *     6000
 * );
 *
 * If the interaction was already acknowledged,
 * this automatically falls back to an ephemeral follow-up.
 */

async function temporaryEphemeralReply(
    interaction,
    options,
    deleteAfter =
        6000
) {
    if (
        interaction.replied ||
        interaction.deferred
    ) {
        return temporaryEphemeralFollowUp(
            interaction,
            options,
            deleteAfter
        );
    }


    await interaction.reply({
        ...options,

        flags:
            MessageFlags.Ephemeral
    });


    scheduleDeleteReply(
        interaction,
        deleteAfter
    );
}


/*
 * ============================================================
 * TEMPORARY EPHEMERAL FOLLOW-UP
 * ============================================================
 *
 * Primarily used after deferUpdate().
 */

async function temporaryEphemeralFollowUp(
    interaction,
    options,
    deleteAfter =
        6000
) {
    const message =
        await interaction.followUp({
            ...options,

            flags:
                MessageFlags.Ephemeral,

            fetchReply:
                true
        });


    if (
        deleteAfter >
        0
    ) {
        setTimeout(
            async () => {
                try {
                    await interaction.webhook.deleteMessage(
                        message.id
                    );
                } catch {}
            },
            deleteAfter
        );
    }


    return message;
}


/*
 * ============================================================
 * TEMPORARY REPLY DELETE
 * ============================================================
 */

function scheduleDeleteReply(
    interaction,
    deleteAfter
) {
    if (
        !deleteAfter ||
        deleteAfter <=
        0
    ) {
        return;
    }


    setTimeout(
        async () => {
            try {
                await interaction.deleteReply();
            } catch {}
        },
        deleteAfter
    );
}


/*
 * ============================================================
 * RANDOM INTEGER
 * ============================================================
 *
 * Inclusive:
 *
 * randomInteger(1, 10)
 *
 * can return 1 through 10.
 */

function randomInteger(
    min,
    max
) {
    const minimum =
        Math.ceil(
            min
        );


    const maximum =
        Math.floor(
            max
        );


    return Math.floor(
        Math.random() *
        (
            maximum -
            minimum +
            1
        )
    )
    +
    minimum;
}


/*
 * ============================================================
 * RANDOM FLOAT
 * ============================================================
 */

function randomFloat(
    min,
    max
) {
    return (
        Math.random() *
        (
            max -
            min
        )
    )
    +
    min;
}


/*
 * ============================================================
 * RANDOM ITEM
 * ============================================================
 */

function randomItem(
    items
) {
    if (
        !Array.isArray(
            items
        )
        ||
        items.length ===
            0
    ) {
        return undefined;
    }


    return items[
        Math.floor(
            Math.random() *
            items.length
        )
    ];
}


/*
 * ============================================================
 * FORMAT PERCENTAGE
 * ============================================================
 *
 * formatPercentage(0.523)
 *
 * => "52.3%"
 */

function formatPercentage(
    value,
    decimals =
        1
) {
    return (
        (
            value *
            100
        )
            .toFixed(
                decimals
            )
        +
        "%"
    );
}


/*
 * ============================================================
 * VALIDATION
 * ============================================================
 */

function validateCooldownName(
    cooldownName
) {
    if (
        typeof cooldownName !==
            "string"
        ||
        !/^[A-Za-z0-9_]+$/.test(
            cooldownName
        )
    ) {
        throw new Error(
            "INVALID_COOLDOWN_NAME"
        );
    }
}


function validateCooldownDuration(
    durationSeconds
) {
    if (
        !Number.isSafeInteger(
            durationSeconds
        )
        ||
        durationSeconds <
            0
    ) {
        throw new Error(
            "INVALID_COOLDOWN_DURATION"
        );
    }
}


/*
 * ============================================================
 * DATABASE INDEXES
 * ============================================================
 *
 * Call this once during startup if desired.
 *
 * The unique cooldown index is particularly useful for
 * claimCooldown() because it prevents duplicate
 * userId + guildId cooldown records.
 */

async function ensureCookieIndexes() {
    const {
        cookies,
        cooldowns
    } =
        getCollections();


    await cookies.createIndex(
        {
            userId:
                1
        },
        {
            unique:
                true
        }
    );


    await cooldowns.createIndex(
        {
            userId:
                1,

            guildId:
                1
        },
        {
            unique:
                true
        }
    );
}


/*
 * ============================================================
 * EXPORTS
 * ============================================================
 */

module.exports = {
    getCollections,

    getCookieBalance,
    changeCookies,
    changeCookiesOnce,
    hasCookieOperation,
    clearCookieOperation,
    transferCookies,

    getCooldown,
    setCooldown,
    claimCooldown,
    clearCooldown,

    temporaryEphemeralReply,
    temporaryEphemeralFollowUp,

    getUnixTime,

    randomInteger,
    randomFloat,
    randomItem,
    formatPercentage,

    ensureCookieIndexes
};