conjurer/librarian_commands.py

import asyncio
import io
import logging
import os
import random
import time
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,
    mark_delivered,
    submit_ai_query,
)
from constants import (
    DIR_PATH_SADOX,
    LIBRARIAN_SERVICE_ADDRESS,
    QUERY_STATUS,
    SELF_CALLBACK,
    SEND_QUERY,
    service_headers,
)
from librarian_watchdog import FLAG, pending_verdict

SERVICE_HEADERS = service_headers()

# Per-query watchdog tuning.
PENDING_WATCH_SECONDS = 30            # how often to ask the librarian about a uuid
PENDING_GRACE_SECONDS = 45           # unknown-but-pending must persist this long
PENDING_HARD_TTL = 60 * 60 * 24 * 3  # drop tracking after 3 days no matter what


class DataModule(commands.Cog):
    def __init__(self, bot, logger_name):
        self.bot = bot
        self.logger = logging.getLogger(logger_name)
        # uuid -> {"ctx", "query", "created", "unknown_since"} for every search
        # dispatched but not yet answered. watch_pending polls the librarian for
        # each; check_data_q removes an entry the moment its result is rendered.
        self.pending = {}

    def _track_pending(self, query_uuid, query, ctx):
        """Start watching a dispatched search so a lost result can be caught."""
        self.pending[str(query_uuid)] = {
            "ctx": ctx,
            "query": query,
            "created": time.monotonic(),
            "unknown_since": None,
        }

    @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)
            if not res:
                await ctx.send("*Conjurer grzebie w pustej skrzyni* Nie ma dziś żadnych komiksów.")
                return
            # randrange(len) is 0..len-1; the old randrange(0, len-1) never picked
            # the last file and raised ValueError('empty range') on a single file.
            # trunk-ignore(bandit/B311)
            filename = res[random.randrange(len(res))]
            # 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:
                # The result arrived and is about to be rendered - stop the
                # watchdog from ever flagging this uuid as lost.
                self.pending.pop(str(fresh_data.uuid), None)
                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.red/{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 = ""

                # The result is now on screen: mark it delivered so the
                # librarian's resends become no-ops and it is dropped from the
                # durable inbox (never replayed again). Done after the core
                # render but before the optional AI review, which is a bonus.
                mark_delivered(str(fresh_data.uuid))

                # 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

    @tasks.loop(seconds=PENDING_WATCH_SECONDS)
    async def watch_pending(self):
        """Per-query safety net for lost results (case a).

        For each dispatched-but-unanswered search, ask the librarian whether it
        still knows the uuid (queued or processing). While it does, the search is
        progressing - leave it alone (a busy librarian is fine). The moment a
        uuid VANISHES on the librarian while still pending here, its result was
        computed but never reached us: after a short grace window (to rule out a
        result that is merely in flight) we tell the channel - but ONLY then.

        A normally-delivered result is popped from self.pending by check_data_q,
        so it never reaches the flag path.
        """
        now = time.monotonic()
        for query_uuid in list(self.pending.keys()):
            info = self.pending.get(query_uuid)
            if info is None:
                continue
            # Hard cap so a permanently-unreachable librarian can't leak entries.
            if now - info["created"] > PENDING_HARD_TTL:
                self.logger.warning("Dropping stale pending query %s (hard TTL)", query_uuid)
                self.pending.pop(query_uuid, None)
                continue
            try:
                response = await asyncio.to_thread(
                    requests.post,
                    f"{LIBRARIAN_SERVICE_ADDRESS}{QUERY_STATUS}",
                    json={"UUID": query_uuid},
                    headers=SERVICE_HEADERS,
                    timeout=5,
                )
                known = (
                    response.status_code == 200
                    and response.json().get("data", {}).get("known", False)
                )
            except (
                requests.exceptions.RequestException,
                ValueError,
                AttributeError,
                KeyError,
                TypeError,
            ) as exc:
                # Librarian unreachable / garbled or unexpected answer: we can't
                # judge, so don't cry wolf, and don't let one bad poll kill the
                # loop. Reset the clock and try again next tick.
                self.logger.info("Pending check for %s inconclusive: %s", query_uuid, exc)
                info["unknown_since"] = None
                continue
            action, info["unknown_since"] = pending_verdict(
                known, info["unknown_since"], now, PENDING_GRACE_SECONDS
            )
            # Re-check membership: the await above yields, so check_data_q may
            # have just delivered (and popped) this result.
            if action == FLAG and query_uuid in self.pending:
                await self._flag_lost(query_uuid, info)
                self.pending.pop(query_uuid, None)

    async def _flag_lost(self, query_uuid, info):
        """Tell the querent their finished search never made it back."""
        message = (
            "*Winda na książki z hukiem wraca z podziemi PUSTA. Z głośnika trzeszczy:* "
            f"Twoje zapytanie {query_uuid} (\"{info['query']}\") przemieliło się w "
            "bibliotece do końca, ale wynik przepadł gdzieś w drodze do baru - nic nie "
            "dotarło. Zawołaj szefa albo puść jeszcze raz."
        )
        try:
            await info["ctx"].send(message)
        except Exception:  # pylint: disable=broad-exception-caught
            self.logger.exception("Failed to post lost-result notice for %s", query_uuid)

    @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,
            "callback": SELF_CALLBACK,
        }
        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)
        self._track_pending(query_uuid, query, ctx)
        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,
            "callback": SELF_CALLBACK,
        }
        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)
        self._track_pending(query_uuid, query, ctx)
        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,
            "callback": SELF_CALLBACK,
        }
        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)
        self._track_pending(query_uuid, query, ctx)
        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()
    dm.watch_pending.start()
    await bot.add_cog(dm)
    logger.info("Loading data sharing commands module done")