From d98d4156fe1e0bcf6f9f3bc4e4de2fe666232d4e Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Thu, 2 Mar 2023 18:45:56 +0530
Subject: [PATCH] feat: utility functions for poll updates

---
 src/Utils/messages.ts | 68 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)

diff --git a/src/Utils/messages.ts b/src/Utils/messages.ts
index bf1624b..7a4ef9a 100644
--- a/src/Utils/messages.ts
+++ b/src/Utils/messages.ts
@@ -24,6 +24,7 @@ import {
 	WATextMessage,
 } from '../Types'
 import { isJidGroup, jidNormalizedUser } from '../WABinary'
+import { sha256 } from './crypto'
 import { generateMessageID, getKeyAuthor, unixTimestampSeconds } from './generics'
 import { downloadContentFromMessage, encryptedStream, generateThumbnail, getAudioDuration, MediaDownloadOptions } from './messages-media'
 
@@ -706,6 +707,73 @@ export const updateMessageWithReaction = (msg: Pick<WAMessage, 'reactions'>, rea
 	msg.reactions = reactions
 }
 
+/** Update the message with a new poll update */
+export const updateMessageWithPollUpdate = (
+	msg: Pick<WAMessage, 'pollUpdates'>,
+	update: proto.IPollUpdate
+) => {
+	const authorID = getKeyAuthor(update.pollUpdateMessageKey)
+
+	const reactions = (msg.pollUpdates || [])
+		.filter(r => getKeyAuthor(r.pollUpdateMessageKey) !== authorID)
+	if(update.vote?.selectedOptions?.length) {
+		reactions.push(update)
+	}
+
+	msg.pollUpdates = reactions
+}
+
+type VoteAggregation = {
+	name: string
+	voters: string[]
+}
+
+/**
+ * Aggregates all poll updates in a poll.
+ * @param msg the poll creation message
+ * @param meId your jid
+ * @returns A list of options & their voters
+ */
+export function getAggregateVotesInPollMessage(
+	{ message, pollUpdates }: Pick<WAMessage, 'pollUpdates' | 'message'>,
+	meId?: string
+) {
+	const opts = message?.pollCreationMessage?.options || []
+	const voteHashMap = opts.reduce((acc, opt) => {
+		const hash = sha256(Buffer.from(opt.optionName || '')).toString()
+		acc[hash] = {
+			name: opt.optionName || '',
+			voters: []
+		}
+		return acc
+	}, {} as { [_: string]: VoteAggregation })
+
+	for(const update of pollUpdates || []) {
+		const { vote } = update
+		if(!vote) {
+			continue
+		}
+
+		for(const option of vote.selectedOptions || []) {
+			const hash = option.toString()
+			let data = voteHashMap[hash]
+			if(!data) {
+				voteHashMap[hash] = {
+					name: 'Unknown',
+					voters: []
+				}
+				data = voteHashMap[hash]
+			}
+
+			voteHashMap[hash].voters.push(
+				getKeyAuthor(update.pollUpdateMessageKey, meId)
+			)
+		}
+	}
+
+	return Object.values(voteHashMap)
+}
+
 /** Given a list of message keys, aggregates them by chat & sender. Useful for sending read receipts in bulk */
 export const aggregateMessageKeysNotFromMe = (keys: proto.IMessageKey[]) => {
 	const keyMap: { [id: string]: { jid: string, participant: string | undefined, messageIds: string[] } } = { }