recalculate_referral_rewards.mjs

import admin from "firebase-admin";
import fs from "fs";
import pLimit from "p-limit";
import * as XLSX from "xlsx";

const serviceAccount = JSON.parse(
  fs.readFileSync("firebase_staging_creds.json", "utf8")
);

admin.initializeApp({
  credential: admin.credential.cert(serviceAccount),
});

const db = admin.firestore();
const limit = pLimit(400); // Increased concurrency limit
const checkpointFile = "checkpoint.json";
const logFile = "recalculated_users.xlsx"; // Log file for testing
const testArr = []; // Array to store new users
const maxUsersToProcess = 10000; // Maximum number of users to process

function saveCheckpoint(lastDocId) {
  fs.writeFileSync(checkpointFile, JSON.stringify({ lastDocId }), "utf8");
}

function loadCheckpoint() {
  if (fs.existsSync(checkpointFile)) {
    const data = fs.readFileSync(checkpointFile, "utf8");
    return JSON.parse(data).lastDocId;
  }
  return null;
}

async function getUserRewardsCreatedInSeasonTwo(userId) {
  try {
    const seasonTwoLaunchDate = new Date("2024-06-17");

    const snapshot = await db
      .collection("rewardsV2")
      .where("userId", "==", userId)
      .where("createdAt", ">=", seasonTwoLaunchDate)
      .get();

    // Remove rewards that have type as recurringReferral.
    // A where clause could also be added to the query above
    // to filter out these rewards but it would require an index to be created
    // so we're filtering them out after getting the rewards.
    const nonReferralRewards = snapshot.docs.filter(
      (doc) => doc.data().type !== "recurringReferral"
    );

    const rewards = [];
    nonReferralRewards.forEach((doc) => rewards.push(doc.data()));

    return rewards;
  } catch (error) {
    console.error(`Error fetching rewards for user ${userId}:`, error);
    throw error;
  }
}

async function calculateReferralRewards(referrerUserId) {
  if (!referrerUserId) {
    return { reward: 0, referralCount: 0, totalReferralRewards: 0 }; // No referral rewards if referrerUserId is undefined
  }

  try {
    const snapshot = await db
      .collection("referrals")
      .where("referrerUserId", "==", referrerUserId)
      .get();

    let totalPoints = 0;
    let referralCount = 0;
    console.log(`Referrer ${referrerUserId} has ${snapshot.size} referrals.`); // Debug log

    const referralPromises = snapshot.docs.map(async (doc) => {
      const data = doc.data();
      const referredUserId = data.referredUserId;
      referralCount++;
      console.log(
        `Processing referral ${referralCount} for referrer ${referrerUserId}: referred user ${referredUserId}`
      ); // Debug log

      try {
        const referredUserRewards = await getUserRewardsCreatedInSeasonTwo(
          referredUserId
        );
        const referredUserTotalPoints = referredUserRewards.reduce(
          (sum, reward) => sum + (parseInt(reward.points, 10) || 0),
          0
        );

        totalPoints += referredUserTotalPoints;
      } catch (error) {
        console.error(
          `Error calculating rewards for referred user ${referredUserId}:`,
          error
        );
      }
    });

    await Promise.all(referralPromises);

    const reward = totalPoints * 0.2;
    return { reward, referralCount };
  } catch (error) {
    console.error(
      `Error calculating referral rewards for referrer ${referrerUserId}:`,
      error
    );
    throw error;
  }
}

async function setUserPoints(userId, newPoints, oldPoints) {
  try {
    await db.collection("users").doc(userId).update({
      points: newPoints,
      oldPoints: oldPoints,
    });
  } catch (error) {
    console.error(`Error setting points for user ${userId}:`, error);
    throw error;
  }
}

async function recalculateUserPoints(userId, referrerUserId, eoaAddress) {
  try {
    const userDoc = await db.collection("users").doc(userId).get();
    const oldPoints = userDoc.data().points || 0;

    const userRewards = await getUserRewardsCreatedInSeasonTwo(userId);
    const totalUserRewards = userRewards.reduce(
      (sum, reward) => sum + reward.points,
      0
    );

    const { reward: referralReward, referralCount } =
      await calculateReferralRewards(referrerUserId);
    const newTotalPoints = totalUserRewards + referralReward;

    await setUserPoints(userId, newTotalPoints, oldPoints);
    console.log(`User ${userId} new total points: ${newTotalPoints}`);

    // Log the recalculated user details to the test array
    testArr.push({
      eoaAddress,
      oldPoints,
      newPoints: newTotalPoints,
      referralCount,
    });
  } catch (error) {
    console.error(`Error recalculating points for user ${userId}:`, error);
  }
}

async function recalculateAllUsersPoints(batchSize = 800) {
  let lastDocId = loadCheckpoint();
  let lastDoc = null;
  let usersProcessed = 0;
  let hasMoreUsers = true;

  if (lastDocId) {
    lastDoc = await db.collection("users").doc(lastDocId).get();
  }

  do {
    try {
      let query = db
        .collection("users")
        .orderBy("points", "desc")
        .limit(Math.min(batchSize, maxUsersToProcess - usersProcessed));
      if (lastDoc) {
        query = query.startAfter(lastDoc);
      }

      const usersSnapshot = await query.get();
      if (usersSnapshot.empty) {
        hasMoreUsers = false;
        break;
      }

      const users = [];
      usersSnapshot.forEach((doc) => {
        const userData = doc.data();
        users.push({
          id: doc.id,
          eoaAddress: userData.eoaAddress,
        });
      });

      await Promise.all(
        users.map((user) =>
          limit(() => recalculateUserPoints(user.id, user.id, user.eoaAddress))
        )
      );

      lastDoc = usersSnapshot.docs[usersSnapshot.docs.length - 1];
      lastDocId = lastDoc.id;
      console.log(checkpointFile);
      saveCheckpoint(lastDocId);

      usersProcessed += users.length;
      console.log(`Processed ${usersProcessed} users so far...`);

      if (usersProcessed >= maxUsersToProcess) {
        hasMoreUsers = false;
        break;
      }
    } catch (error) {
      console.error("Error processing batch of users:", error);
    }
  } while (hasMoreUsers);

  console.log("Recalculated points for all users.");

  // Write the test array to an .xlsx file
  const worksheet = XLSX.utils.json_to_sheet(testArr);
  const workbook = XLSX.utils.book_new();
  XLSX.utils.book_append_sheet(workbook, worksheet, "Recalculated Users");
  XLSX.writeFile(workbook, logFile);

  // Optionally, remove the checkpoint file if processing is complete
  fs.unlinkSync(checkpointFile);
}

recalculateAllUsersPoints()
  .catch(console.error)
  .finally(() => {
    console.log("Test Array:", testArr); // Print the test array at the end
  });

Refactored to query rewards created on season 2.