From ce019dcf0f0e3c3f1a8c57b776aa5fe39854b598 Mon Sep 17 00:00:00 2001
From: Rafael Garcia <rafaeljosegarciasuarez@gmail.com>
Date: Sun, 14 May 2023 12:50:45 -0500
Subject: [PATCH] feat(store/labels): add ability to store and query labels

---
 src/Store/make-in-memory-store.ts | 296 +++++++++++++++++++++---------
 src/Store/object-repository.ts    |  31 ++++
 2 files changed, 243 insertions(+), 84 deletions(-)
 create mode 100644 src/Store/object-repository.ts

diff --git a/src/Store/make-in-memory-store.ts b/src/Store/make-in-memory-store.ts
index f8d7a6b..5787b24 100644
--- a/src/Store/make-in-memory-store.ts
+++ b/src/Store/make-in-memory-store.ts
@@ -1,46 +1,97 @@
-import type KeyedDB from '@adiwajshing/keyed-db'
+import KeyedDB from '@adiwajshing/keyed-db'
 import type { Comparable } from '@adiwajshing/keyed-db/lib/Types'
 import type { Logger } from 'pino'
 import { proto } from '../../WAProto'
 import { DEFAULT_CONNECTION_CONFIG } from '../Defaults'
 import type makeMDSocket from '../Socket'
 import type { BaileysEventEmitter, Chat, ConnectionState, Contact, GroupMetadata, PresenceData, WAMessage, WAMessageCursor, WAMessageKey } from '../Types'
+import { Label } from '../Types/Label'
+import { LabelAssociation, LabelAssociationType, MessageLabelAssociation } from '../Types/LabelAssociation'
 import { toNumber, updateMessageWithReaction, updateMessageWithReceipt } from '../Utils'
 import { jidNormalizedUser } from '../WABinary'
 import makeOrderedDictionary from './make-ordered-dictionary'
+import { ObjectRepository } from './object-repository'
 
 type WASocket = ReturnType<typeof makeMDSocket>
 
 export const waChatKey = (pin: boolean) => ({
 	key: (c: Chat) => (pin ? (c.pinned ? '1' : '0') : '') + (c.archived ? '0' : '1') + (c.conversationTimestamp ? c.conversationTimestamp.toString(16).padStart(8, '0') : '') + c.id,
-	compare: (k1: string, k2: string) => k2.localeCompare (k1)
+	compare: (k1: string, k2: string) => k2.localeCompare(k1)
 })
 
 export const waMessageID = (m: WAMessage) => m.key.id || ''
 
+export const waLabelAssociationKey: Comparable<LabelAssociation, string> = {
+	key: (la: LabelAssociation) => (la.type === LabelAssociationType.Chat ? la.chatId + la.labelId : la.chatId + la.messageId + la.labelId),
+	compare: (k1: string, k2: string) => k2.localeCompare(k1)
+}
+
 export type BaileysInMemoryStoreConfig = {
 	chatKey?: Comparable<Chat, string>
+	labelAssociationKey?: Comparable<LabelAssociation, string>
 	logger?: Logger
 }
 
 const makeMessagesDictionary = () => makeOrderedDictionary(waMessageID)
 
-export default (
-	{ logger: _logger, chatKey }: BaileysInMemoryStoreConfig
-) => {
-	const logger = _logger || DEFAULT_CONNECTION_CONFIG.logger.child({ stream: 'in-mem-store' })
-	chatKey = chatKey || waChatKey(true)
-	const KeyedDB = require('@adiwajshing/keyed-db').default as new (...args: any[]) => KeyedDB<Chat, string>
+const predefinedLabels = Object.freeze<Record<string, Label>>({
+	'0': {
+		id: '0',
+		name: 'New customer',
+		predefinedId: '0',
+		color: 0,
+		deleted: false
+	},
+	'1': {
+		id: '1',
+		name: 'New order',
+		predefinedId: '1',
+		color: 1,
+		deleted: false
+	},
+	'2': {
+		id: '2',
+		name: 'Pending payment',
+		predefinedId: '2',
+		color: 2,
+		deleted: false
+	},
+	'3': {
+		id: '3',
+		name: 'Paid',
+		predefinedId: '3',
+		color: 3,
+		deleted: false
+	},
+	'4': {
+		id: '4',
+		name: 'Order completed',
+		predefinedId: '4',
+		color: 4,
+		deleted: false
+	}
+})
 
-	const chats = new KeyedDB(chatKey, c => c.id)
-	const messages: { [_: string]: ReturnType<typeof makeMessagesDictionary> } = { }
-	const contacts: { [_: string]: Contact } = { }
-	const groupMetadata: { [_: string]: GroupMetadata } = { }
-	const presences: { [id: string]: { [participant: string]: PresenceData } } = { }
+export default (
+	{ logger: _logger, chatKey, labelAssociationKey }: BaileysInMemoryStoreConfig
+) => {
+	// const logger = _logger || DEFAULT_CONNECTION_CONFIG.logger.child({ stream: 'in-mem-store' })
+	chatKey = chatKey || waChatKey(true)
+	labelAssociationKey = labelAssociationKey || waLabelAssociationKey
+	const logger = _logger || DEFAULT_CONNECTION_CONFIG.logger.child({ stream: 'in-mem-store' })
+	// const KeyedDB = require('@adiwajshing/keyed-db').default as new (...args: any[]) => KeyedDB<Chat, string>
+
+	const chats = new KeyedDB<Chat, string>(chatKey, c => c.id)
+	const messages: { [_: string]: ReturnType<typeof makeMessagesDictionary> } = {}
+	const contacts: { [_: string]: Contact } = {}
+	const groupMetadata: { [_: string]: GroupMetadata } = {}
+	const presences: { [id: string]: { [participant: string]: PresenceData } } = {}
 	const state: ConnectionState = { connection: 'close' }
+	const labels = new ObjectRepository<Label>(predefinedLabels)
+	const labelAssociations = new KeyedDB<LabelAssociation, string>(labelAssociationKey, labelAssociationKey.key)
 
 	const assertMessageList = (jid: string) => {
-		if(!messages[jid]) {
+		if (!messages[jid]) {
 			messages[jid] = makeMessagesDictionary()
 		}
 
@@ -49,7 +100,7 @@ export default (
 
 	const contactsUpsert = (newContacts: Contact[]) => {
 		const oldContacts = new Set(Object.keys(contacts))
-		for(const contact of newContacts) {
+		for (const contact of newContacts) {
 			oldContacts.delete(contact.id)
 			contacts[contact.id] = Object.assign(
 				contacts[contact.id] || {},
@@ -60,6 +111,12 @@ export default (
 		return oldContacts
 	}
 
+	const labelsUpsert = (newLabels: Label[]) => {
+		for (const label of newLabels) {
+			labels.upsertById(label.id, label)
+		}
+	}
+
 	/**
 	 * binds to a BaileysEventEmitter.
 	 * It listens to all events and constructs a state that you can query accurate data from.
@@ -77,10 +134,10 @@ export default (
 			messages: newMessages,
 			isLatest
 		}) => {
-			if(isLatest) {
+			if (isLatest) {
 				chats.clear()
 
-				for(const id in messages) {
+				for (const id in messages) {
 					delete messages[id]
 				}
 			}
@@ -89,13 +146,13 @@ export default (
 			logger.debug({ chatsAdded }, 'synced chats')
 
 			const oldContacts = contactsUpsert(newContacts)
-			for(const jid of oldContacts) {
+			for (const jid of oldContacts) {
 				delete contacts[jid]
 			}
 
 			logger.debug({ deletedContacts: oldContacts.size, newContacts }, 'synced contacts')
 
-			for(const msg of newMessages) {
+			for (const msg of newMessages) {
 				const jid = msg.key.remoteJid!
 				const list = assertMessageList(jid)
 				list.upsert(msg, 'prepend')
@@ -105,8 +162,8 @@ export default (
 		})
 
 		ev.on('contacts.update', updates => {
-			for(const update of updates) {
-				if(contacts[update.id!]) {
+			for (const update of updates) {
+				if (contacts[update.id!]) {
 					Object.assign(contacts[update.id!], update)
 				} else {
 					logger.debug({ update }, 'got update for non-existant contact')
@@ -117,71 +174,98 @@ export default (
 			chats.upsert(...newChats)
 		})
 		ev.on('chats.update', updates => {
-			for(let update of updates) {
+			for (let update of updates) {
 				const result = chats.update(update.id!, chat => {
-					if(update.unreadCount! > 0) {
+					if (update.unreadCount! > 0) {
 						update = { ...update }
 						update.unreadCount = (chat.unreadCount || 0) + update.unreadCount!
 					}
 
 					Object.assign(chat, update)
 				})
-				if(!result) {
+				if (!result) {
 					logger.debug({ update }, 'got update for non-existant chat')
 				}
 			}
 		})
+
+		ev.on('labels.edit', (label: Label) => {
+			if (label.deleted) {
+				return labels.deleteById(label.id)
+			}
+
+			// WhatsApp can store only up to 20 labels
+			if (labels.count() < 20) {
+				return labels.upsertById(label.id, label)
+			}
+
+			logger.error('Labels count exceed')
+		})
+
+		ev.on('labels.association', ({ type, association }) => {
+			switch (type) {
+				case 'add':
+					labelAssociations.upsert(association)
+					break
+				case 'remove':
+					labelAssociations.delete(association)
+					break
+				default:
+					console.error(`unknown operation type [${type}]`)
+			}
+		})
+
 		ev.on('presence.update', ({ id, presences: update }) => {
 			presences[id] = presences[id] || {}
 			Object.assign(presences[id], update)
 		})
 		ev.on('chats.delete', deletions => {
-			for(const item of deletions) {
+			for (const item of deletions) {
 				chats.deleteById(item)
 			}
 		})
 		ev.on('messages.upsert', ({ messages: newMessages, type }) => {
 			switch (type) {
-			case 'append':
-			case 'notify':
-				for(const msg of newMessages) {
-					const jid = jidNormalizedUser(msg.key.remoteJid!)
-					const list = assertMessageList(jid)
-					list.upsert(msg, 'append')
+				case 'append':
+				case 'notify':
+					for (const msg of newMessages) {
+						const jid = jidNormalizedUser(msg.key.remoteJid!)
+						const list = assertMessageList(jid)
+						list.upsert(msg, 'append')
 
-					if(type === 'notify') {
-						if(!chats.get(jid)) {
-							ev.emit('chats.upsert', [
-								{
-									id: jid,
-									conversationTimestamp: toNumber(msg.messageTimestamp),
-									unreadCount: 1
-								}
-							])
+						if (type === 'notify') {
+							if (!chats.get(jid)) {
+								ev.emit('chats.upsert', [
+									{
+										id: jid,
+										conversationTimestamp: toNumber(msg.messageTimestamp),
+										unreadCount: 1
+									}
+								])
+							}
 						}
 					}
-				}
 
-				break
+					break
 			}
 		})
 		ev.on('messages.update', updates => {
-			for(const { update, key } of updates) {
+			for (const { update, key } of updates) {
 				const list = assertMessageList(key.remoteJid!)
 				const result = list.updateAssign(key.id!, update)
-				if(!result) {
+				if (!result) {
 					logger.debug({ update }, 'got update for non-existent message')
 				}
 			}
 		})
 		ev.on('messages.delete', item => {
-			if('all' in item) {
+			if ('all' in item) {
 				const list = messages[item.jid]
 				list?.clear()
 			} else {
 				const jid = item.keys[0].remoteJid!
 				const list = messages[jid]
-				if(list) {
+				if (list) {
 					const idSet = new Set(item.keys.map(k => k.id))
 					list.filter(m => !idSet.has(m.key.id))
 				}
@@ -189,9 +273,9 @@ export default (
 		})
 
 		ev.on('groups.update', updates => {
-			for(const update of updates) {
+			for (const update of updates) {
 				const id = update.id!
-				if(groupMetadata[id]) {
+				if (groupMetadata[id]) {
 					Object.assign(groupMetadata[id], update)
 				} else {
 					logger.debug({ update }, 'got update for non-existant group metadata')
@@ -201,42 +285,42 @@ export default (
 
 		ev.on('group-participants.update', ({ id, participants, action }) => {
 			const metadata = groupMetadata[id]
-			if(metadata) {
+			if (metadata) {
 				switch (action) {
-				case 'add':
-					metadata.participants.push(...participants.map(id => ({ id, isAdmin: false, isSuperAdmin: false })))
-					break
-				case 'demote':
-				case 'promote':
-					for(const participant of metadata.participants) {
-						if(participants.includes(participant.id)) {
-							participant.isAdmin = action === 'promote'
+					case 'add':
+						metadata.participants.push(...participants.map(id => ({ id, isAdmin: false, isSuperAdmin: false })))
+						break
+					case 'demote':
+					case 'promote':
+						for (const participant of metadata.participants) {
+							if (participants.includes(participant.id)) {
+								participant.isAdmin = action === 'promote'
+							}
 						}
-					}
 
-					break
-				case 'remove':
-					metadata.participants = metadata.participants.filter(p => !participants.includes(p.id))
-					break
+						break
+					case 'remove':
+						metadata.participants = metadata.participants.filter(p => !participants.includes(p.id))
+						break
 				}
 			}
 		})
 
 		ev.on('message-receipt.update', updates => {
-			for(const { key, receipt } of updates) {
+			for (const { key, receipt } of updates) {
 				const obj = messages[key.remoteJid!]
 				const msg = obj?.get(key.id!)
-				if(msg) {
+				if (msg) {
 					updateMessageWithReceipt(msg, receipt)
 				}
 			}
 		})
 
 		ev.on('messages.reaction', (reactions) => {
-			for(const { key, reaction } of reactions) {
+			for (const { key, reaction } of reactions) {
 				const obj = messages[key.remoteJid!]
 				const msg = obj?.get(key.id!)
-				if(msg) {
+				if (msg) {
 					updateMessageWithReaction(msg, reaction)
 				}
 			}
@@ -246,15 +330,25 @@ export default (
 	const toJSON = () => ({
 		chats,
 		contacts,
-		messages
+		messages,
+		labels,
+		labelAssociations
 	})
 
-	const fromJSON = (json: { chats: Chat[], contacts: { [id: string]: Contact }, messages: { [id: string]: WAMessage[] } }) => {
+	const fromJSON = (json: {
+		chats: Chat[],
+		contacts: { [id: string]: Contact },
+		messages: { [id: string]: WAMessage[] },
+		labels: { [labelId: string]: Label },
+		labelAssociations: LabelAssociation[]
+	}) => {
 		chats.upsert(...json.chats)
+		labelAssociations.upsert(...json.labelAssociations || [])
 		contactsUpsert(Object.values(json.contacts))
-		for(const jid in json.messages) {
+		labelsUpsert(Object.values(json.labels || {}))
+		for (const jid in json.messages) {
 			const list = assertMessageList(jid)
-			for(const msg of json.messages[jid]) {
+			for (const msg of json.messages[jid]) {
 				list.upsert(proto.WebMessageInfo.fromObject(msg), 'append')
 			}
 		}
@@ -268,17 +362,19 @@ export default (
 		groupMetadata,
 		state,
 		presences,
+		labels,
+		labelAssociations,
 		bind,
 		/** loads messages from the store, if not found -- uses the legacy connection */
-		loadMessages: async(jid: string, count: number, cursor: WAMessageCursor) => {
+		loadMessages: async (jid: string, count: number, cursor: WAMessageCursor) => {
 			const list = assertMessageList(jid)
 			const mode = !cursor || 'before' in cursor ? 'before' : 'after'
 			const cursorKey = !!cursor ? ('before' in cursor ? cursor.before : cursor.after) : undefined
 			const cursorValue = cursorKey ? list.get(cursorKey.id!) : undefined
 
 			let messages: WAMessage[]
-			if(list && mode === 'before' && (!cursorKey || cursorValue)) {
-				if(cursorValue) {
+			if (list && mode === 'before' && (!cursorKey || cursorValue)) {
+				if (cursorValue) {
 					const msgIdx = list.array.findIndex(m => m.key.id === cursorKey?.id)
 					messages = list.array.slice(0, msgIdx)
 				} else {
@@ -286,7 +382,7 @@ export default (
 				}
 
 				const diff = count - messages.length
-				if(diff < 0) {
+				if (diff < 0) {
 					messages = messages.slice(-count) // get the last X messages
 				}
 			} else {
@@ -295,27 +391,59 @@ export default (
 
 			return messages
 		},
-		loadMessage: async(jid: string, id: string) => messages[jid]?.get(id),
-		mostRecentMessage: async(jid: string) => {
+		/**
+		 * Get all available labels for profile
+		 *
+		 * Keep in mind that the list is formed from predefined tags and tags
+		 * that were "caught" during their editing.
+		 */
+		getLabels: () => {
+			return labels
+		},
+
+		/**
+		 * Get labels for chat
+		 *
+		 * @returns Label IDs
+		 **/
+		getChatLabels: (chatId: string) => {
+			return labelAssociations.filter((la) => la.chatId === chatId).all()
+		},
+
+		/**
+		 * Get labels for message
+		 *
+		 * @returns Label IDs
+		 **/
+		getMessageLabels: (messageId: string) => {
+			const associations = labelAssociations
+				.filter((la: MessageLabelAssociation) => la.messageId === messageId)
+				.all()
+
+			return associations.map(({ labelId }) => labelId)
+
+		},
+		loadMessage: async (jid: string, id: string) => messages[jid]?.get(id),
+		mostRecentMessage: async (jid: string) => {
 			const message: WAMessage | undefined = messages[jid]?.array.slice(-1)[0]
 			return message
 		},
-		fetchImageUrl: async(jid: string, sock: WASocket | undefined) => {
+		fetchImageUrl: async (jid: string, sock: WASocket | undefined) => {
 			const contact = contacts[jid]
-			if(!contact) {
+			if (!contact) {
 				return sock?.profilePictureUrl(jid)
 			}
 
-			if(typeof contact.imgUrl === 'undefined') {
+			if (typeof contact.imgUrl === 'undefined') {
 				contact.imgUrl = await sock?.profilePictureUrl(jid)
 			}
 
 			return contact.imgUrl
 		},
-		fetchGroupMetadata: async(jid: string, sock: WASocket | undefined) => {
-			if(!groupMetadata[jid]) {
+		fetchGroupMetadata: async (jid: string, sock: WASocket | undefined) => {
+			if (!groupMetadata[jid]) {
 				const metadata = await sock?.groupMetadata(jid)
-				if(metadata) {
+				if (metadata) {
 					groupMetadata[jid] = metadata
 				}
 			}
@@ -332,7 +460,7 @@ export default (
 
 		// 	return groupMetadata[jid]
 		// },
-		fetchMessageReceipts: async({ remoteJid, id }: WAMessageKey) => {
+		fetchMessageReceipts: async ({ remoteJid, id }: WAMessageKey) => {
 			const list = messages[remoteJid!]
 			const msg = list?.get(id!)
 			return msg?.userReceipt
@@ -347,7 +475,7 @@ export default (
 		readFromFile: (path: string) => {
 			// require fs here so that in case "fs" is not available -- the app does not crash
 			const { readFileSync, existsSync } = require('fs')
-			if(existsSync(path)) {
+			if (existsSync(path)) {
 				logger.debug({ path }, 'reading from file')
 				const jsonStr = readFileSync(path, { encoding: 'utf-8' })
 				const json = JSON.parse(jsonStr)
diff --git a/src/Store/object-repository.ts b/src/Store/object-repository.ts
new file mode 100644
index 0000000..9cd0ec1
--- /dev/null
+++ b/src/Store/object-repository.ts
@@ -0,0 +1,31 @@
+export class ObjectRepository<T extends object> {
+    readonly entityMap: Map<string, T>
+
+    constructor(entities: Record<string, T> = {}) {
+        this.entityMap = new Map(Object.entries(entities))
+    }
+
+    findById(id: string) {
+        return this.entityMap.get(id)
+    }
+
+    findAll() {
+        return Array.from(this.entityMap.values())
+    }
+
+    upsertById(id: string, entity: T) {
+        return this.entityMap.set(id, { ...entity })
+    }
+
+    deleteById(id: string) {
+        return this.entityMap.delete(id)
+    }
+
+    count() {
+        return this.entityMap.size
+    }
+
+    toJSON() {
+        return this.findAll()
+    }
+}
\ No newline at end of file