dakiny-bot/cogs/clan_messages/__init__.py

import os

import nextcord
import uvicorn
from fastapi import FastAPI
from git import Repo
from nextcord.ext import commands
from sqlalchemy import select, update
from sqlalchemy.orm import Session

from config import DakinyConfig
from database_setup.sql_data_classes import Messages, engine

app = FastAPI()
REPO_URL = "https://gitea.cheezers.org/imcompiling/dakiny-messages.git"
REPO_DIR = "./dakiny-messages"
MESSAGES_DIR = f"{REPO_DIR}/messages"


class ClanMessages(commands.Cog, name="Clan Messages"):
    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.config: DakinyConfig = self.bot.config
        self.repo_setup()

    @commands.Cog.listener()
    async def on_ready(self):
        """
        Triggers when the bot is in ready state. Adds the webhook listener to FastAPI and serves.
        """
        await self.update_messages()
        app.add_api_route("/message/", self.hello, methods=["POST"])
        config = uvicorn.Config(app=app, port=5000, log_level="info")
        server = uvicorn.Server(config)
        await server.serve()

    async def hello(self):
        await self.update_messages()
        return

    async def update_messages(self):
        """
        _summary_
        """

        self.origin.pull()
        guild: nextcord.Guild = self.bot.get_guild(self.config.params.guild_id)

        # for ever directory, sort the files by number then post
        # posting involves checking if there is a message for that particular file then editing it or creating a
        # new one. the message could be deleted, which would cause a new message to be sent and database updated.
        dirs = [
            dir_name
            for dir_name in os.listdir(MESSAGES_DIR)
            if os.path.isdir(f"{MESSAGES_DIR}/{dir_name}")
        ]

        with Session(engine) as session:
            for dir_name in dirs:
                dir_path = f"{MESSAGES_DIR}/{dir_name}"
                channel = await self.get_clan_message_channel(session, dir_name, guild)
                sorted_files = sorted(
                    os.listdir(dir_path),
                    key=lambda filename: int(filename.split(".")[0]),
                )

                for file_name in sorted_files:
                    message_name = f"{dir_name}/{file_name}"
                    file_path = f"{dir_path}/{file_name}"
                    with open(file_path, "r", encoding="utf8") as reader:
                        stmt = select(Messages.Id).where(Messages.Path == message_name)
                        msg_id = session.scalar(stmt)
                        if msg_id is None:
                            # don't have a message for this file
                            msg = await channel.send(reader.read(2000))
                            new_message = Messages(Path=message_name, Id=msg.id)
                            session.add(new_message)
                            session.commit()
                        else:
                            try:
                                msg: nextcord.Message = await channel.fetch_message(
                                    msg_id
                                )
                                await msg.edit(reader.read(2000))
                            except nextcord.NotFound:
                                msg = await channel.send(reader.read(2000))
                                stmt = (
                                    update(Messages)
                                    .where(Messages.Path == message_name)
                                    .values(Id=msg.id)
                                )
                                session.execute(stmt)
                                session.commit()

    async def get_clan_message_channel(
        self, session: Session, channel_name: str, guild: nextcord.Guild
    ) -> nextcord.TextChannel:
        """
        Used to obtain a channel from the guild object. Searches database for an id to get the channel. If that is not
        found, then searches through all channels in the guild to find one with a matching name. If that also doesn't
        turn up a channel, then creates it. Adds an entry into database for last two paths.
        Assumes that there is an open session to the database.

        Args:
            session (Session): Open connection to database.
            channel_name (str): Name of channel to obtain.
            guild (nextcord.Guild): Guild object to get channel from.

        Returns:
            nextcord.TextChannel: Channel object with name channel_name.
        """
        stmt = select(Messages.Id).where(Messages.Path == channel_name)
        channel_id = session.scalar(stmt)
        if channel_id is None:
            ret_channel = None
            # channel is not tracked. see if it exists
            channels = await guild.fetch_channels()
            for channel in channels:
                if (
                    channel.name == channel_name
                    and channel.__class__ is nextcord.TextChannel
                ):
                    ret_channel = channel

            if ret_channel is None:
                # channel doesn't exist, we must create it
                ret_channel = await guild.create_text_channel(
                    channel_name,
                    reason="Creating channel as required for by design of Dakiny-Messages repo.",
                )

            # record channel information
            new_item = Messages(Path=channel_name, Id=ret_channel.id)
            session.add(new_item)
            session.commit()
            return ret_channel

        return await guild.fetch_channel(channel_id)

    def repo_setup(self):
        """
        Checks if the repo exists and creates repo object, clones if it does not.
        """
        # check if repo has been cloned. open repo object
        if os.path.isdir(REPO_DIR):
            self.repo = Repo(REPO_DIR)
        else:
            self.repo = Repo.clone_from(REPO_URL, REPO_DIR)

        self.origin = self.repo.remote("origin")

        if self.config.params.git_branch == "test":
            self.repo.create_head("test", self.origin.refs.test).set_tracking_branch(
                self.origin.refs.test
            ).checkout()
        else:
            self.repo.create_head("main", self.origin.refs.main).set_tracking_branch(
                self.origin.refs.main
            ).checkout()

        self.origin.pull()


# This function will be called when this extension is loaded.
# It is necessary to add these functions to the bot.
def setup(bot: commands.Bot):
    """Adds function to bot"""
    bot.add_cog(ClanMessages(bot))