diff --git a/database/schema.sql b/database/schema.sql index e69de29..22c9782 100644 --- a/database/schema.sql +++ b/database/schema.sql @@ -0,0 +1,12 @@ +CREATE TABLE IF NOT EXISTS starboard_entries ( + msg_id BIGINT PRIMARY KEY, + bot_message_id BIGINT, + channel BIGINT, + stars INT NOT NULL DEFAULT 1, + bot_content_id BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS starers ( + user_id BIGINT, + msg_id BIGINT +) \ No newline at end of file diff --git a/modules/stars.py b/modules/stars.py index 7bf102b..b58f897 100644 --- a/modules/stars.py +++ b/modules/stars.py @@ -20,8 +20,267 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ + +from __future__ import annotations +import asyncpg +import discord +import datetime +from discord.ext import commands + + import core +CONFIG = core.CONFIG["STARBOARD"] +STARBOARD_CHANNEL_ID = CONFIG.get("starboard_channel_id") + +STARBOARD_EMBED_COLOR = 0xFFFF00 +STARBOARD_EMOJI = "⭐" +HEADER_TEMPLATE = "**{}** {} in: <#{}> ID: {}" + +VALID_FILE_ATTACHMENTS = (".jpg", ".jpeg", ".png", ".webp", ".gif") +VIDEO_FILE_ATTACHMENTS = (".mp4", ".mov") + + +class JumpView(discord.ui.View): + def __init__( + self, + *, + timeout: float, + url: str | None, + label_name: str = "Jump to message", + ) -> None: + super().__init__(timeout=timeout) + self.add_item(discord.ui.Button(url=url, label=label_name, style=discord.ButtonStyle.primary)) + + +class StarboardEntry: + exists: bool = False + msg_id: int = 0 + channel_id: int = 0 + stars: int = 0 + bot_message_id: int = 0 + bot_content_id: int = 0 + + db: asyncpg.Pool[asyncpg.Record] + + def __init__(self, db: asyncpg.Pool[asyncpg.Record], msg_id: int) -> None: + self.msg_id = msg_id + self.db: asyncpg.Pool[asyncpg.Record] = db + + @classmethod + async def from_query(cls, db: asyncpg.Pool[asyncpg.Record], msg_id: int): + instance = cls(db, msg_id) + await instance.fetch() + return instance + + async def fetch(self) -> None: + query = """SELECT * FROM starboard_entries WHERE msg_id=$1""" + + result = await self.db.fetchrow(query, self.msg_id) + + if result is None: + self.exists = False + return + + self.exists = True + self.msg_id = result["msg_id"] + self.channel_id = result["channel"] + self.stars = result["stars"] + self.bot_message_id = result["bot_message_id"] + self.bot_content_id = result["bot_content_id"] + + +class Starboard(core.Cog): + def __init__(self, bot: core.Bot) -> None: + self.bot = bot + + self.remove_on_delete: bool = CONFIG.get("remove_on_delete") + self.entry_requirement: int = CONFIG.get("entry_requirement") + self.starboard_channel_id: int = CONFIG.get("starboard_channel_id") + self.pool: asyncpg.Pool[asyncpg.Record] = bot.pool + + def get_star(self, stars: int) -> str: + if stars <= 2: + return "⭐" + elif stars <= 4: + return "🌟" + elif stars <= 6: + return "💫" + else: + return "✨" + + async def add_entry( + self, message_id: int, bot_message_id: int, payload_channel_id: int, reactions: int, content_id: int + ) -> None: + query = """INSERT INTO starboard_entries VALUES ( + $1, + $2, + $3, + $4, + $5 + )""" + await self.pool.execute(query, message_id, bot_message_id, payload_channel_id, reactions, content_id) + + async def add_starer(self, user_id: int, message_id: int) -> None: + query = """ + INSERT INTO starers VALUES ( + $1, + $2 + )""" + + await self.pool.execute(query, user_id, message_id) + + async def remove_starer(self, message_id: int, user_id: int) -> None: + query = """DELETE FROM starers WHERE msg_id = $1 AND user_id= $2""" + await self.pool.execute(query, message_id, user_id) + + async def update_entry(self, reactions: int, message_id: int) -> None: + query = """UPDATE starboard_entries SET stars = $1 WHERE msg_id = $2""" + await self.pool.execute(query, reactions, message_id) + + async def remove_entry(self, message_id: int) -> None: + query = """DELETE FROM starboard_entries WHERE msg_id= $1""" + await self.pool.execute(query, message_id) + + async def clear_starers(self, message_id: int) -> None: + query = """DELETE FROM starers WHERE msg_id = $1""" + await self.pool.execute(query, message_id) + + def get_formatted_time(self) -> str: + now = datetime.datetime.now() + time = now.strftime("%m/%d/%Y %I:%M %p") + return time + + async def handle_star(self, payload: discord.RawReactionActionEvent) -> None: + time = self.get_formatted_time() + entry = await StarboardEntry.from_query(self.pool, payload.message_id) + + if str(payload.emoji) != STARBOARD_EMOJI: + return + + channel: discord.TextChannel = self.bot.get_channel(payload.channel_id) # type: ignore + message: discord.Message = await channel.fetch_message(payload.message_id) + + reaction = discord.utils.get(message.reactions, emoji=STARBOARD_EMOJI) + reaction_count = reaction.count if reaction else 0 + + if entry.exists: + bot_msg_id = entry.bot_message_id + + query = """SELECT * FROM starers WHERE user_id=$1 AND msg_id=$2""" + + starer = await self.pool.fetchrow(query, payload.user_id, entry.msg_id) + + if starer is not None: + return + + await self.add_starer(payload.user_id, payload.message_id) + + bot_channel: discord.TextChannel = self.bot.get_channel(self.starboard_channel_id) # type: ignore + bot_message = await bot_channel.fetch_message(bot_msg_id) + + stars = reaction_count + star = self.get_star(stars) + await bot_message.edit(content=HEADER_TEMPLATE.format(star, stars, payload.channel_id, payload.channel_id)) + await self.update_entry(stars, payload.message_id) + return + + if reaction_count < self.entry_requirement: + return + + star = self.get_star(reaction_count) + + embed = discord.Embed(color=STARBOARD_EMBED_COLOR, description=message.content) + if len(message.attachments) > 0: + for attachment in message.attachments: + filename = attachment.filename + if filename.endswith(VALID_FILE_ATTACHMENTS): + if attachment.is_spoiler(): + embed.add_field(name="", value=f"[Click to view spoiler]({attachment.url})", inline=True) + continue + embed.set_image(url=attachment.url) + elif filename.endswith(VIDEO_FILE_ATTACHMENTS): + embed.add_field(name="", value=f"[File: {attachment.filename}]({message.jump_url})") + else: + continue + message_url: str = message.jump_url + + embed.set_author(name=message.author.display_name, icon_url=message.author.avatar) + embed.set_footer(text=time) + + starboard = self.bot.get_channel(self.starboard_channel_id) + + bot_message: discord.Message = await starboard.send( # type: ignore + HEADER_TEMPLATE.format(star, reaction_count, payload.channel_id, payload.channel_id) + ) + content_message = await starboard.send( # type: ignore + embed=embed, + view=JumpView(url=message_url, timeout=40), + ) + + await self.add_entry(message.id, bot_message.id, payload.channel_id, reaction.count, content_message.id) + await self.add_starer(payload.user_id, message.id) + + async def handle_unstar(self, payload: discord.RawReactionActionEvent) -> None: + entry = await StarboardEntry.from_query(self.pool, payload.message_id) + + bot_msg_id = entry.bot_message_id + content_id = entry.bot_content_id + + if not entry.exists: + return + + channel: discord.TextChannel = await self.bot.fetch_channel(self.starboard_channel_id) # type: ignore + bot_msg = await channel.fetch_message(bot_msg_id) + content_msg = await channel.fetch_message(content_id) + + reacted_message_channel: discord.TextChannel = await self.bot.fetch_channel(payload.channel_id) # type: ignore + reacted_message = await reacted_message_channel.fetch_message(payload.message_id) + + reaction: discord.Reaction | None = discord.utils.get(reacted_message.reactions, emoji=STARBOARD_EMOJI) + reaction_count: int = reaction.count if reaction else 0 + if reaction_count == 0: + # not possible to have zero stars. + await bot_msg.delete() + await content_msg.delete() + + await self.remove_entry(payload.message_id) + return + + star = self.get_star(reaction_count) + message = HEADER_TEMPLATE.format(star, reaction_count, payload.channel_id, payload.channel_id) + + await self.update_entry(reaction_count, payload.message_id) + await self.remove_starer(payload.message_id, payload.user_id) + await bot_msg.edit(content=message) + + @core.Cog.listener() + async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent) -> None: + await self.handle_star(payload) + + @commands.Cog.listener() + async def on_raw_reaction_remove(self, payload: discord.RawReactionActionEvent) -> None: + await self.handle_unstar(payload) + + @commands.Cog.listener() + async def on_message_delete(self, message: discord.Message) -> None: + possible_entry = await StarboardEntry.from_query(self.pool, message.id) # type: ignore + await possible_entry.fetch() + if not possible_entry.exists: + return + + if not self.remove_on_delete: + return + + channel: discord.TextChannel = await self.bot.fetch_channel(self.starboard_channel_id) # type: ignore + bot_msg = await channel.fetch_message(possible_entry.bot_message_id) + content_msg = await channel.fetch_message(possible_entry.bot_content_id) + + await bot_msg.delete() + await content_msg.delete() + await self.remove_entry(message.id) + await self.clear_starers(message.id) + async def setup(bot: core.Bot) -> None: - pass + await bot.add_cog(Starboard(bot)) diff --git a/types_/config.py b/types_/config.py index c2b4fde..8401efa 100644 --- a/types_/config.py +++ b/types_/config.py @@ -34,6 +34,12 @@ class Suggestions(TypedDict): webhook_url: str +class Starboard(TypedDict): + remove_on_delete: bool + entry_requirement: int + starboard_channel_id: int + + class Config(TypedDict): prefix: str owner_ids: NotRequired[list[int]] @@ -43,3 +49,4 @@ class Config(TypedDict): SNEKBOX: NotRequired[Snekbox] BADBIN: BadBin SUGGESTIONS: NotRequired[Suggestions] + STARBOARD: Starboard