From 33a5a503d9f704509cf47e2904e8c39c43b2fb0d Mon Sep 17 00:00:00 2001
From: Pedro Howat <pedro@instant.com.br>
Date: Tue, 3 Sep 2024 14:20:37 -0300
Subject: [PATCH] Adds option to fetch instance usage data

---
 .../kwik/controllers/kwik.controller.ts       | 47 ++++++++++++++++++-
 .../integrations/kwik/routes/kwik.router.ts   | 20 +++++++-
 src/api/routes/instance.router.ts             | 15 +++++-
 3 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/src/api/integrations/kwik/controllers/kwik.controller.ts b/src/api/integrations/kwik/controllers/kwik.controller.ts
index 501fe0b2..f9589320 100644
--- a/src/api/integrations/kwik/controllers/kwik.controller.ts
+++ b/src/api/integrations/kwik/controllers/kwik.controller.ts
@@ -1,4 +1,4 @@
-import { Document } from 'bson';
+import { calculateObjectSize, Document } from 'bson';
 
 import { configService, Database } from '../../../../config/env.config';
 import { Logger } from '../../../../config/logger.config';
@@ -137,4 +137,49 @@ export class KwikController {
 
     return { status: 'ok' };
   }
+  public async instanceInfo({ instanceName }: InstanceDto, messageTimestamp: number, usage?: number) {
+    const db = configService.get<Database>('DATABASE');
+    const connection = dbserver.getClient().db(db.CONNECTION.DB_PREFIX_NAME + '-whatsapp-api');
+    const messages = connection.collection('messages');
+    const pipeline: Document[] = [
+      { $sort: { 'key.remoteJid': -1, messageTimestamp: -1 } },
+      {
+        $group: {
+          _id: '$key.remoteJid',
+          owner: { $first: '$owner' },
+          message: { $first: '$message' },
+          lastAllMsgTimestamp: { $first: '$messageTimestamp' },
+          name: { $first: '$pushName' },
+          fromMe: { $first: '$key.fromMe' },
+        },
+      },
+      { $match: { owner: instanceName, lastAllMsgTimestamp: { $gte: messageTimestamp } } },
+      { $count: 'rowCount' },
+    ];
+    const chatCount = await messages.aggregate(pipeline).toArray();
+
+    if (usage) {
+      return {
+        chatCount: chatCount[0].rowCount,
+        totalSize: usage,
+        newVal: 0,
+      };
+    } else {
+      const userMessages = await messages
+        .find({ owner: instanceName, messageTimestamp: { $gte: messageTimestamp } })
+        .toArray();
+
+      let totalSize = 0;
+
+      userMessages.forEach(function (doc) {
+        totalSize += calculateObjectSize(doc);
+      });
+
+      return {
+        chatCount: chatCount[0].rowCount,
+        totalSize: totalSize,
+        newVal: 1,
+      };
+    }
+  }
 }
diff --git a/src/api/integrations/kwik/routes/kwik.router.ts b/src/api/integrations/kwik/routes/kwik.router.ts
index f57d3cc4..effa2b2e 100644
--- a/src/api/integrations/kwik/routes/kwik.router.ts
+++ b/src/api/integrations/kwik/routes/kwik.router.ts
@@ -36,7 +36,7 @@ export class KwikRouter extends RouterBroker {
       return res.status(HttpStatus.OK).json(response);
     });
     this.router.post(this.routerPath('cleanup'), ...guards, async (req, res) => {
-      logger.verbose('request received in findChats');
+      logger.verbose('request received in cleanup');
       logger.verbose('request body: ');
       logger.verbose(req.body);
 
@@ -52,6 +52,24 @@ export class KwikRouter extends RouterBroker {
 
       return res.status(HttpStatus.OK).json(response);
     });
+
+    this.router.get(this.routerPath('instanceInfo'), ...guards, async (req, res) => {
+      logger.verbose('request received in instanceInfo');
+      logger.verbose('request body: ');
+      logger.verbose(req.body);
+
+      logger.verbose('request query: ');
+      logger.verbose(req.query);
+
+      const response = await this.dataValidate<InstanceDto>({
+        request: req,
+        schema: null,
+        ClassRef: InstanceDto,
+        execute: (instance) => kwikController.instanceInfo(instance, Number(req.query.messageTimestamp)),
+      });
+
+      return res.status(HttpStatus.OK).json(response);
+    });
   }
 
   public readonly router = Router();
diff --git a/src/api/routes/instance.router.ts b/src/api/routes/instance.router.ts
index d2cc6fb8..a85df612 100644
--- a/src/api/routes/instance.router.ts
+++ b/src/api/routes/instance.router.ts
@@ -6,7 +6,7 @@ import { dbserver } from '../../libs/db.connect';
 import { instanceNameSchema, oldTokenSchema, presenceOnlySchema } from '../../validate/validate.schema';
 import { RouterBroker } from '../abstract/abstract.router';
 import { InstanceDto, SetPresenceDto } from '../dto/instance.dto';
-import { instanceController } from '../server.module';
+import { instanceController, kwikController } from '../server.module';
 import { OldToken } from '../services/auth.service';
 import { HttpStatus } from './index.router';
 
@@ -104,6 +104,9 @@ export class InstanceRouter extends RouterBroker {
         logger.verbose(req.body);
 
         const key = req.get('apikey');
+        const fullFetch = req.query.fullFetch ? Number(req.query.fullFetch) : 0;
+        const messageTimestamp = req.query.messageTimestamp ? req.query.messageTimestamp : -1;
+        const usage = req.query.usage ? Number(req.query.usage) : null;
 
         logger.verbose('request query: ');
         logger.verbose(req.query);
@@ -114,6 +117,16 @@ export class InstanceRouter extends RouterBroker {
           execute: (instance) => instanceController.fetchInstances(instance, key),
         });
 
+        if (fullFetch == 1) {
+          const usageData = await this.dataValidate<InstanceDto>({
+            request: req,
+            schema: null,
+            ClassRef: InstanceDto,
+            execute: (instance) => kwikController.instanceInfo(instance, Number(messageTimestamp), usage),
+          });
+          response.usageData = usageData;
+        }
+
         return res.status(HttpStatus.OK).json(response);
       })
       .post(this.routerPath('setPresence'), ...guards, async (req, res) => {