pixivcord/src/main.py

import asyncio
import logging
import os
from shlex import join
from tempfile import TemporaryDirectory

import discord

import discord_tools.message
import discord_tools.pixiv_link_extractor
import pixiv.pixiv_api as pixiv_api

logging.basicConfig(
    format="[%(asctime)s] [%(levelname)8s] %(name)s: %(message)s", level=logging.INFO
)


REFRESH_TOKEN = os.getenv("PIXIV_REFRESH_TOKEN")
if REFRESH_TOKEN is None:
    logging.fatal("Pixiv refresh token is missing, please set PIXIV_REFRESH_TOKEN")
    exit(1)
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
if DISCORD_TOKEN is None:
    logging.fatal("Discord token is missing, please set DISCORD_TOKEN")
    exit(1)


intents = discord.Intents.default()
intents.message_content = True


class MyClient(discord.Client):
    def __init__(self, *, intents: discord.Intents):
        super().__init__(intents=intents)
        self.tree = discord.app_commands.CommandTree(self)

    async def setup_hook(self):
        await self.tree.sync()


client = MyClient(intents=intents)


@client.event
async def on_ready():
    logging.info(f"We have logged in as {client.user}")


@client.tree.command()
@discord.app_commands.describe(
    url="Pixiv art url",
)
async def pixivprev(interaction: discord.Interaction, url: str):
    """Create pixiv preview"""
    (result, match) = discord_tools.pixiv_link_extractor.check_pixiv_url(url)
    if result and match is not None:
        await interaction.response.defer()
        with TemporaryDirectory() as tmp_dir:
            loop = asyncio.get_running_loop()
            data = await loop.run_in_executor(
                None,
                pixiv_api.download_pixiv_images,
                int(match.group("id")),
                tmp_dir,
                REFRESH_TOKEN,
            )
            if data is None:
                await interaction.response.send_message(
                    "Fail to extract pixiv image", ephemeral=True
                )
                return
            (title, author) = data
            files = os.listdir(tmp_dir)
            files.sort()
            await discord_tools.message.send_command_reply(
                interaction,
                files,
                tmp_dir,
                True,
                title=title,
                art_author=author,
                url=url,
                author=interaction.user,
            )
    else:
        await interaction.response.send_message("Invalid pixiv url", ephemeral=True)


@client.tree.context_menu(name="Create pixiv preview")
async def pixiv_preview(interaction: discord.Interaction, message: discord.Message):
    (result, muted, match) = discord_tools.pixiv_link_extractor.should_process(message)
    if result and match is not None:
        await interaction.response.defer()
        with TemporaryDirectory() as tmp_dir:
            loop = asyncio.get_running_loop()
            data = await loop.run_in_executor(
                None,
                pixiv_api.download_pixiv_images,
                int(match.group("id")),
                tmp_dir,
                REFRESH_TOKEN,
            )
            if data is None:
                await interaction.response.send_message(
                    "Fail to extract pixiv image", ephemeral=True
                )
                return
            files = os.listdir(tmp_dir)
            files.sort()
            (title, author) = data
            await discord_tools.message.send_command_reply(
                interaction=interaction,
                files=files,
                tmp_dir=tmp_dir,
                muted=muted,
                title=title,
                art_author=author,
                url=match.group("url"),
                author=message.author,
            )
    else:
        await interaction.response.send_message(
            "Can't find any valid pixiv link.", ephemeral=True
        )


@client.event
async def on_guild_join(guild: discord.Guild):
    logging.info(f"Joined new guild: {guild.name}")


@client.event
async def on_message(message: discord.Message):
    if message.author == client.user:
        return
    (result, muted, match) = discord_tools.pixiv_link_extractor.should_process(message)
    if result and match is not None:
        with TemporaryDirectory() as tmp_dir:
            async with message.channel.typing():
                loop = asyncio.get_running_loop()
                data = await loop.run_in_executor(
                    None,
                    pixiv_api.download_pixiv_images,
                    int(match.group("id")),
                    tmp_dir,
                    REFRESH_TOKEN,
                )
                if data is None:
                    return
                files = os.listdir(tmp_dir)
                files.sort()
                # await send_message.send_message_with_embed(
                #     message, files, title, tmp_dir, match.group("url")
                # )
                await discord_tools.message.send_message(message, files, tmp_dir, muted)


client.run(DISCORD_TOKEN)