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, submit_ai_query
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
                # ai_lines is a clean, plain rendering of the SAME list in the
                # SAME (Crossref-relevance) order, for the optional AI review.
                ai_lines = []
                l_p = 1
                for doi in fresh_data.entries:
                    self.logger.info(doi)
                    desc = fresh_data.entries[doi]
                    title = desc["Title"][0] if desc.get("Title") else "(bez tytułu)"
                    entries.append(
                        f"{l_p}. {title} pod linkiem https://www.sci-hub.se/{doi} i jest to {desc['type']}\n"
                    )
                    ai_lines.append(f"{l_p}. {title} (DOI: {doi}, typ: {desc['type']})")
                    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 = ""

                # Optional AI pass: re-rank the (already Crossref-relevance-
                # sorted) DOI list and review the sources. Enqueued to the AI
                # worker so it runs on whatever backend $gadaj_teraz selected;
                # the answer lands in this same channel.
                if getattr(fresh_data, "ai_review", False) and ai_lines:
                    target = getattr(ctx, "channel", ctx)
                    review_prompt = (
                        f'Poniżej lista źródeł naukowych znalezionych dla zapytania: "{query}".\n'
                        "Lista jest już wstępnie posortowana według trafności wg Crossref "
                        "(od najtrafniejszej).\n\n"
                        "Twoje zadania:\n"
                        "1. Przeważ i uporządkuj listę według RZECZYWISTEJ trafności do zapytania "
                        "(najtrafniejsze u góry).\n"
                        "2. Do każdej pozycji dopisz jedno-, dwuzdaniową recenzję: typ i wiarygodność "
                        "źródła oraz dlaczego (nie) pasuje do zapytania.\n"
                        "Odpowiedz zwięźle, numerowaną listą, po polsku.\n\n"
                        "Źródła:\n" + "\n".join(ai_lines)
                    )
                    submit_ai_query(
                        prompt=review_prompt,
                        channel_id=target.id,
                        request_type="NONE",
                        username=searcher,
                        query_uuid=str(fresh_data.uuid),
                    )
                    await ctx.send(
                        "*Conjurer podaje listę naszemu rezydentowi-mądrali od AI* "
                        "Za chwilę dorzuci recenzję i swoje przesortowanie wg trafności."
                    )

                # 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="wyszukaj_z_recenzja",
        description="Jak wyszukaj_linki_do_dokumentow, ale wyniki przesortuje trafnością i zrecenzuje AI",
        guild=discord.Object(id=664789470779932693),
    )
    async def wyszukaj_z_recenzja(self, ctx):
        """Same as wyszukaj_linki_do_dokumentow, but flags the search for an AI
        review: when the DOI hits come back, the list + the search phrase are sent
        to the bot's AI backend for a weighted-relevance re-rank and a source
        review, delivered to this channel. The flag rides on the QueryControl so
        it survives the round-trip and is matched back to this search by UUID.
        """
        query = ctx.message.content
        query_uuid = uuid.uuid4()
        ctx.message.content = ctx.message.content.replace("$wyszukaj_z_recenzja", "")

        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 notuje, wrzuca liścik do rury pneumatycznej i mruży oko* Tym razem jak coś"
            + " znajdę, przepuszczę wyniki jeszcze przez naszego rezydenta-mądralę od AI - przeważy"
            + " trafność i zrecenzuje źródła. Poczekaj kilka godzin - biblioteka to 3/4 stacji."
        )
        query_response = await coroutine
        if not query_response.status_code == 200:
            await ctx.send(
                "*Z rury wydobywa się dym. Conjurer pryska w nią pierwszą cieczą pod ręką i wybucha"
                + " drobny pożar.* Wołaj szefa - mam 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, ai_review=True)
        OUT_COMM_Q.put(query_object)
        await ctx.send(
            f"Poszło z recenzją AI. Identyfikator: {query_uuid}. Jesteś {queue_size} w kolejce."
            + " Najpierw dojadą surowe wyniki, a zaraz po nich przesortowanie i recenzja od AI."
        )

    @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")