Heimdallr/Heimdallr.py

import logging
from os import getenv

from naff import (
    Client,
    Intents,
    OptionTypes,
    Permissions,
    listen,
    slash_command,
    InteractionContext,
    slash_option,
)
from naff.api import events
from naff.models.discord.embed import (
    Embed,
    EmbedAuthor,
    EmbedField,
    EmbedFooter,
)


from database import GuildSettings as GuildSettingsModel, JoinLeave

from dotenv import load_dotenv

bot = Client(intents=Intents.ALL, debug_scope=387153131378835456)


@listen()
async def on_ready():
    print(f"Bot '{bot.user.username}' is ready!")
    print(f"This bot is owned by {bot.owner}")

    for guild in bot.guilds:
        guild_q = GuildSettingsModel.get_or_none(GuildSettingsModel.guild_id == guild.id)
        if guild_q is None:
            guild_q = GuildSettingsModel(guild_id=guild.id)
            guild_q.save()
        


# @listen()
# async def on_message_create(event: MessageCreate):
#     print(f"{event.message.author.username}: {event.message.content}")

@listen(events.MemberAdd)
async def on_member_join(event: events.MemberAdd):
    joinleave_q: JoinLeave
    joinleave_q, _ = JoinLeave.get_or_create(guild_id=event.guild.id)

    if joinleave_q.message_channel is None:
        return
    channel = await bot.fetch_channel(joinleave_q.message_channel)

    if not joinleave_q.join_message_enabled:
        return

    if joinleave_q.join_message is None or joinleave_q.join_message == "":
        await channel.send(
            f"{event.member.mention} has joined the server!"
        )
        return

    await channel.send(str(joinleave_q.join_message).format(member=event.member, guild=event.guild))

@listen(events.MemberRemove)
async def on_member_leave(event: events.MemberRemove):
    joinleave_q: JoinLeave
    joinleave_q, _ = JoinLeave.get_or_create(guild_id=event.guild.id)

    if joinleave_q.message_channel is None:
        return
    channel = await bot.fetch_channel(joinleave_q.message_channel)

    if not joinleave_q.leave_message_enabled:
        return

    if joinleave_q.leave_message is None or joinleave_q.leave_message == "":
        await channel.send(
            f"{event.member.mention} has left the server!"
        )
        return

    await channel.send(str(joinleave_q.leave_message).format(member=event.member, guild=event.guild))

    


@slash_command(name="ping", description="Ping the bot")
async def ping_command(ctx: InteractionContext):
    ctx.ephemeral = True
    await ctx.send("Pong!")


@slash_command(name="bot-info", description="Get info about the bot")
async def bot_info_command(ctx: InteractionContext):
    await ctx.send(
        ephemeral=True,
        embed=Embed(
            title=f"{bot.user.username}",
            description=f"This bot is owned by {bot.owner}",
            color=0x00FF00,
            timestamp=bot.user.created_at,
            url=bot.user.avatar.as_url(),
            author=EmbedAuthor(
                name=bot.user.username,
                icon_url=bot.user.avatar.as_url(),
            ),
            footer=EmbedFooter(
                text=f"Created at {bot.user.created_at}",
                icon_url=bot.user.avatar.as_url(),
            ),
            fields=[
                EmbedField(
                    name="**Extensions**",
                    value=f"{', '.join(bot.ext.keys())}",
                ),
                EmbedField(
                    name="**Guilds**",
                    value="\n".join([g.name for g in sorted(bot.guilds, key=lambda g: g.name)]),
                ),
            ],
        ),
    )

def set_loglevel(level: str):
    loglevel = logging.WARNING

    match str(level).lower().strip():
        case "d" | "debug":
            loglevel = logging.DEBUG
        case "i" | "info" | "information":
            loglevel = logging.INFO

        case "w" | "warn" | "warning":
            loglevel = logging.WARNING

        case "e" | "error":
            loglevel = logging.ERROR

        case "c" | "critical":
            loglevel = logging.CRITICAL

        case _:
            loglevel = logging.WARNING

    logging.basicConfig(level=loglevel)


def main():
    load_dotenv() # Load environment variables from .env file
    set_loglevel(getenv("HEIMDALLR_LOGLEVEL"))

    # Create basic tables
    GuildSettingsModel.create_table()
    JoinLeave.create_table()

    # Load extensions
    bot.load_extension("commands.admin")
    bot.load_extension("commands.gatekeep") 
    bot.load_extension("commands.quote")
    bot.load_extension("commands.infractions")
    bot.load_extension("commands.self_roles")
    bot.load_extension("commands.polls")
    bot.load_extension("commands.bot_messages")
    bot.start(getenv("DISCORD_TOKEN"))

if __name__ == "__main__":
    main()