conjurer/librarian_commands.py

import asyncio
import io
import logging
import os
import random
import uuid
from queue import Empty

import discord
import pdf2image
import fitz
import PyPDF2
import requests
from discord.ext import commands, tasks

from ai_functions import handle_response
from communication_subroutine import IN_COMM_Q, OUT_COMM_Q, QueryControl
from constants import (
    DIR_PATH_SADOX,
    LIBRARIAN_SERVICE_ADDRESS,
    SEND_QUERY,
    service_headers,
)

SERVICE_HEADERS = service_headers()


class DataModule(commands.Cog):
    def __init__(self, bot, logger_name):
        self.bot = bot
        self.logger = logging.getLogger(logger_name)

    @commands.hybrid_command(
        nsfw=True,
        name="get_image_sadox",
        description="Wyświetla losową stronę z losowego komiksu FanSadox. Bardzo NSFW.",
        guild=discord.Object(id=664789470779932693),
    )
    @commands.has_any_role('Nocna Zmiana', 'Jarl', 'Thane' , 'Bartender')

    async def get_image_sadox(self, ctx):
        """
        Take in a context parameter and retrieve an image related from fansadox collection.

        :param ctx: ctx stands for "context" and is a parameter commonly used in Discord.py commands. It
        represents the context in which the command was invoked, including information such as the message,
        the channel, the server, and the user who invoked the command. In this case, it is likely being used
        to determine
        """
        self.logger.info("Get sadox")
        channel = ctx.message.channel
        async with channel.typing():
            # select random file
            res = []
            # Iterate directory
            for path in os.listdir(DIR_PATH_SADOX):
                # check if current path is a file
                if os.path.isfile(os.path.join(DIR_PATH_SADOX, path)):
                    res.append(path)
            # trunk-ignore(bandit/B311)
            filename = res[random.randrange(0, len(res) - 1)]
            # select random page
            file = open(DIR_PATH_SADOX + filename, "rb")
            if  True:
                doc = fitz.open(DIR_PATH_SADOX + filename)
                totalpages = len(doc)
                # trunk-ignore(bandit/B311)
                page_index = random.randrange(0, totalpages)
                page = doc.load_page(page_index)
                mat = fitz.Matrix(2.0, 2.0)  # powiększenie
                pix = page.get_pixmap(matrix=mat, alpha=False)

                byte_io_stream = io.BytesIO(pix.tobytes("png"))
                byte_io_stream.seek(0)
                byte_io_stream.name = "image.png"
                await ctx.send(file=discord.File(byte_io_stream))
            else: #legacy
                readpdf = PyPDF2.PdfReader(file)
                totalpages = len(readpdf.pages)
                # trunk-ignore(bandit/B311)
                page = random.randrange(1, totalpages)
                # convert page to image
                image = pdf2image.convert_from_path(
                    DIR_PATH_SADOX + filename, first_page=page, last_page=page
                )
                byte_io_stream = io.BytesIO()
                image[0].save(byte_io_stream, "JPEG")
                byte_io_stream.seek(0)
                byte_io_stream.name = "image.jpg"
                file = discord.File(byte_io_stream)
                await ctx.send(file=file)
            self.logger.info("Get sadox completed")

    @tasks.loop(seconds=3)
    async def check_data_q(self):
        """
        This function checks the data queue for any new entries and processes them accordingly.
        """
        try:
            fresh_data = IN_COMM_Q.get(block=False)
            entries = []
            if fresh_data.stop:
                searcher = fresh_data.author
                query = fresh_data.content
                l_p = 1
                for doi in fresh_data.entries:
                    self.logger.info(doi)
                    desc = fresh_data.entries[doi]
                    title = desc["Title"][0]
                    entries.append(
                        f"{l_p}. {title} pod linkiem https://www.sci-hub.se/{doi} i jest to {desc['type']}\n"
                    )
                    l_p += 1
                message = "*Z podłogi wysuwa się winda na książki*"
                if fresh_data.ctx is not None:
                    ctx = fresh_data.ctx
                    message += f" Dzielny poszukiwaczu @{searcher} w odpowiedzi na twoje zapytanie {query} nr {fresh_data.uuid}"
                else:
                    ctx = self.bot.get_channel(1062047571557744721)
                    message += f" Mam wynik dla zapytania {fresh_data.uuid} ale że Hammer coś grzebał nie pamiętam kto to chciał i co chciał"

                message += "nasza biblioteka służy Ci odpowiedzią. Przy dźwięku fanfar winda się otwiera a w środku "
                if len(entries) < 1:
                    message += "niestety nie ma nic"
                    await ctx.send(message)
                elif len(entries) >= 1 and len(entries) < 5:
                    message += " znajduje się coś:\n"
                    for item in entries:
                        message += item
                    await ctx.send(message)
                else:
                    message += " znajduje się cholernie dużo:\n"
                    for item in entries:
                        message += item
                        if len(message) > 1500:
                            await ctx.send(message)
                            message = ""

                # Kept for sentimental reasons
                # await ctx.send(f"O. A tak będzie wyglądało coś ciekawego w przyszłości: {data}")
        except Empty:
            pass

    @commands.hybrid_command(
        name="wyszukaj_linki_do_dokumentow",
        description="Szuka linkow doi w bazie crossref i podaje linki do scihuba",
        guild=discord.Object(id=664789470779932693),
    )
    async def wyszukaj_linki_do_dokumentow(self, ctx):
        """
        The function `wyszukaj_linki_do_dokumentow` searches for links to documents in a crossref database and provides links to scihub.

        :param ctx: The `ctx` parameter is typically used in Discord.py, a Python library for creating Discord bots.
        It represents the context of the command being executed, including information about the message, the server,
        and the user who invoked the command.
        """
        query = ctx.message.content
        query_uuid = uuid.uuid4()
        # TODO: TESTING ONLY!!
        # query_uuid = uuid.UUID('{12345678-1234-5678-1234-567812345678}')
        ctx.message.content = ctx.message.content.replace(
            "$wyszukaj_linki_do_dokumentow", ""
        )

        json_query = {
            "UUID": str(query_uuid),
            "query": str(query),
            "page": 1,
            "deep_search": False,
        }
        coroutine = asyncio.to_thread(
            requests.post,
            f"{LIBRARIAN_SERVICE_ADDRESS}{SEND_QUERY}",
            json=json_query,
            headers=SERVICE_HEADERS,
            timeout=360,
        )
        await ctx.send(
            "*Conjurer powoli notuje podane przez Ciebie. Podchodzi do sprytnie ukrytej przy barze rury od poczty pneumatycznej i"
            + " wrzuca twój liścik do środka* To teraz trza poczekać kilka godzin. Biblioteka to 3/4 stacji."
            + " A przecież pół stacji to browar, motelik dla zabaw cielesnych i ten Bar. Więc 12/8 stacji teraz będzie ciężko pracować"
        )
        query_response = await coroutine
        if not query_response.status_code == 200:
            await ctx.send(
                "*Z rury wydobywa się dym. Conjurer łapie pierwszą ciecz pod ręką i pryska na rurę. Następuje drobna eksplozja i wszystko zaczyna się palić."
                + " Conjurer jest skonfundowany..* Wołaj szefa - mam nieodparte wrażenie że się coś wyjebało"
            )
            return

        query, query_uuid, queue_size = (
            query_response.json()["data"][0],
            query_response.json()["data"][1],
            query_response.json()["data"][2],
        )
        if ctx.message.author.nick:
            username = ctx.message.author.nick
        else:
            username = ctx.message.author.name
        query_object = QueryControl(username, query_uuid, query, ctx)
        OUT_COMM_Q.put(query_object)
        await ctx.send(
            f"No dobra poszło. Zapamiętaj proszę {query_uuid} - to identyfikator twojego zapytania. Jesteś {queue_size} w kolejce."
            + " Zapytania obsługuje algorytm zasilany czterema chomikami zapierdalającymi w kołowrotku - więc wyniki najwcześniej za kilka godzi - ale mogą być też dni."
        )

    @commands.hybrid_command(
        name="glebokie_gardlo",
        description="Przygotowuje drinka o nazwie głębokie gardło",
        guild=discord.Object(id=664789470779932693),
    )

    async def wyszukaj_linki_do_dokumentow_deep(self, ctx):
        """
        The function `wyszukaj_linki_do_dokumentow_deep` performs a deep search for links to documents in
        a crossref database and provides links to scihub.

        :param ctx: The `ctx` parameter is typically used in Discord.py, a Python library for creating Discord bots.
        It represents the context of the command being executed, including information about the message, the server,
        and the user who invoked the command.
        """
        # TODO: Implement deep search logic here
        query = ctx.message.content.replace("$glebokie_gardlo", "")
        allowed = False
        for role in ctx.message.author.roles:
            if role.name == "Bartender":
                allowed = True
            if role.name == "Scribe":
                allowed = True
            if role.name == "Thane":
                allowed = True
        if not allowed:
            if ctx.message.author.nick:
                username = ctx.message.author.nick
            else:
                username = ctx.message.author.name
            vykidailo = False
            bartender = False
            prompt = "Przygotuj mi drinka o nazwie Głębokie Gardło, inspirowanego tym sławnym filmem oraz skandalem Watergate"
            for role in ctx.message.author.roles:
                if role.name == "Vykidailo":
                    vykidailo = True
                if role.name == "Bartender":
                    bartender = True
            global MESSAGE_TABLE  # pylint: disable=global-statement

            result, MESSAGE_TABLE = await handle_response(
                prompt, vykidailo, bartender, MESSAGE_TABLE, username, "GENERAL"
            )
            if len(result) < 1500:
                await ctx.send(result)
            else:
                while len(result) > 1500:
                    await ctx.send(result[:1500])
                    result = result[1500:]
            return

        query_uuid = uuid.uuid4()
        # TODO: TESTING ONLY!!
        # query_uuid = uuid.UUID('{12345678-1234-5678-1234-567812345678}')
        json_query = {
            "UUID": str(query_uuid),
            "query": str(query),
            "page": 1,
            "deep_search": True,
        }
        coroutine = asyncio.to_thread(
            requests.post,
            f"{LIBRARIAN_SERVICE_ADDRESS}{SEND_QUERY}",
            json=json_query,
            headers=SERVICE_HEADERS,
            timeout=360,
        )
        await ctx.send(
            "*Conjurer mruga okiem i zamiast drinka wyjmuje dysk usb z terminala. Podchodzi do sprytnie ukrytej przy barze rury od poczty pneumatycznej i"
            + " wrzuca ten dysk do środk* To teraz trza poczekać kilka godzin. TO będzie głębokie wyszukanie"
            + " Rozgląda się dookoła i stawia przed tobą mętnego drinka o zapachu mięty, gwiezdnego pyłu i oleju silnikowego"
        )
        query_response = await coroutine
        if not query_response.status_code == 200:
            await ctx.send(
                "*Z rury wydobywa się dym. Conjurer łapie pierwszą ciecz pod ręką i pryska na rurę. Następuje drobna eksplozja i wszystko zaczyna się palić."
                + " Conjurer jest skonfundowany..* Wołaj szefa - mam nieodparte wrażenie że się coś wyjebało"
            )
            return

        query, query_uuid, queue_size = (
            query_response.json()["data"][0],
            query_response.json()["data"][1],
            query_response.json()["data"][2],
        )
        if ctx.message.author.nick:
            username = ctx.message.author.nick
        else:
            username = ctx.message.author.name
        query_object = QueryControl(username, query_uuid, query, ctx)
        OUT_COMM_Q.put(query_object)
        await ctx.send(
            f"Wypij wypi {query_uuid} - to identyfikator twojego zapytania. Jesteś {queue_size} w kolejce niestety kolejka obowiazuje zawsze."
            + " Zapytania obsługuje algorytm zasilany czterema chomikami zapierdalającymi w kołowrotku - więc wyniki najwcześniej za kilka godzin - ale mogą być też dni. Głębokie zabawy trwają dłużej. *Znowu mruga*"
        )


async def setup(bot):
    logger = logging.getLogger("discord")
    dm = DataModule(bot, "discord")
    dm.check_data_q.start()
    await bot.add_cog(dm)
    logger.info("Loading data sharing commands module done")