conjurer/communication_subroutine.py

import json
import logging
import os
import re
import threading
import time
import uuid as uuidlib
from queue import Empty, Queue
from typing import Optional
from urllib import request as urequest

import requests
from flask import Flask, abort, jsonify, request
from waitress import serve

HOST_ADDRESS = os.getenv("CONJURER_DISCORD_HOST", "192.168.1.92")
PORT_ADDRESS = int(os.getenv("CONJURER_DISCORD_PORT", "5000"))
ICECAST_ADDRESS = os.getenv("CONJURER_ICECAST", "http://192.168.1.12:8000")
API_KEY = os.getenv("CONJURER_API_KEY")
OUT_COMM_Q = Queue()
IN_COMM_Q = Queue()
# AI request queue: prompts to be answered by the bot's own AI backend (whatever
# $gadaj_teraz currently points at). Drained by the AI cog's worker loop, which
# calls handle_response and delivers the answer to the requested channel. Fed
# either over HTTP (POST /ai_query) or in-process via submit_ai_query().
AI_QUERY_Q = Queue()
SRCHTITLE = re.compile(rb"StreamTitle=\\*(?P<title>[^;]*);").search

awaiting_q = []
incoming_q = Queue()
# A health-check ping whose pong never comes back (dead/dropped librarian) would
# otherwise leave its record in awaiting_q forever. scan_incoming sweeps ping
# records older than this. Kept well above the ping timeout so a slow-but-alive
# round-trip is never swept out from under a waiter.
PING_TTL_SECONDS = 30
app = Flask(__name__)


def _authorize_request() -> None:
    """Reject inbound calls lacking the shared key (no-op if key is unset)."""
    if API_KEY and request.headers.get("X-Conjurer-Api-Key") != API_KEY:
        abort(401)
PREPPED_TRACKS = {
    "requests": "",
    "hit": "",
    "all": "",
    "priority": "",
    "jingles": "",
    "now_playing": "",
    "next": "",
    "meta": "",
}
logger = logging.getLogger("discord")


class QueryControl:
    """
    This class `QueryControl` is used to manage queries with information about the author, UUID,
    content, logger, context, and replies.
    """

    def __init__(self, query_author, query_uuid, query_content, ctx, ai_review=False) -> None:
        self.author = query_author
        self.uuid = query_uuid
        self.content = query_content
        self.logger = logging.getLogger("discord")
        self.stop = False
        self.ctx = ctx
        # When True, once the librarian returns hits, the DOI list + the search
        # phrase are sent to the AI backend for a weighted-relevance re-rank and
        # source review (see librarian_commands.check_data_q).
        self.ai_review = ai_review
        self.logger.info(
            f"Created Query control for {self.author}, {self.uuid}: {self.content}"
        )
        self.replies = []


@app.route("/prepped_tracks", methods=["POST"])
def log_radio_tracks():
    _authorize_request()
    app.logger = logging.getLogger("discord")

    app.logger.info(request)
    record = json.loads(request.data)
    app.logger.info(record)
    if "next" in record[0]:
        metadata = id3(ICECAST_ADDRESS)
        PREPPED_TRACKS["now_playing"] = PREPPED_TRACKS["next"]
        if metadata:
            PREPPED_TRACKS["meta"] = (
                metadata["name"]
                + " - "
                + metadata["title"]
                + "("
                + metadata["genre"]
                + ")"
            )
        else:
            PREPPED_TRACKS["meta"] = "Nie znaju"
    PREPPED_TRACKS[record[0]] = record[1]
    app.logger.info("DATA RECEIVED")
    return jsonify("SUCCESS")


@app.route("/conjurer", methods=["POST"])
def answer_external_command():
    """
    The function `answer_external_command` logs the request data, loads the data as JSON, logs the
    record, and then puts the record into an incoming queue before returning a success message.
    :return: The function `answer_external_command()` is returning a JSON response with the message
    "SUCCESS".
    """
    _authorize_request()
    logger = logging.getLogger("discord")
    logger.info(request)
    record = json.loads(request.data)
    logger.info(record)
    logger.info("DATA RECEIVED")
    incoming_q.put(record)
    return jsonify("SUCCESS")


def submit_ai_query(prompt, channel_id=None, request_type="NONE", username="conjurer", query_uuid=None):
    """Queue an AI prompt for the bot to answer with its configured backend.

    In-process entry point (used by the librarian result handler). ``channel_id``
    is the Discord channel the answer should be posted to; ``request_type`` is
    passed through to handle_response ("NONE" keeps it a clean one-shot that does
    not touch conversation memory).
    """
    AI_QUERY_Q.put(
        {
            "uuid": query_uuid,
            "prompt": prompt,
            "channel_id": channel_id,
            "request_type": request_type,
            "username": username,
        }
    )


@app.route("/ai_query", methods=["POST"])
def ai_query():
    """Inbound AI query: queue a prompt to be answered by the bot's AI backend.

    Payload: {"prompt": str, "channel_id": int, "request_type"?: str,
              "username"?: str, "uuid"?: str}. The prompt is queued and answered
      asynchronously by the AI cog's worker; the answer is posted to channel_id.
    """
    _authorize_request()
    logger = logging.getLogger("discord")
    record = json.loads(request.data)
    prompt = record.get("prompt", "")
    if not prompt:
        return jsonify(isError=True, message="missing 'prompt'", statusCode=400, data=[]), 400
    submit_ai_query(
        prompt=prompt,
        channel_id=record.get("channel_id"),
        request_type=record.get("request_type", "NONE"),
        username=record.get("username", "external"),
        query_uuid=record.get("uuid"),
    )
    logger.info("Queued AI query (qsize=%s)", AI_QUERY_Q.qsize())
    return (
        jsonify(isError=False, message="Queued", statusCode=200, data={"qsize": AI_QUERY_Q.qsize()}),
        200,
    )


@app.route("/conjurer", methods=["GET"])
def check_alive():
    """
    The function "check_alive" returns a JSON response with the message "ALIVE" when the "/conjurer"
    route is accessed via a GET request.
    :return: The code snippet is a Flask route that listens for GET requests to the "/conjurer"
    endpoint. When a GET request is made to this endpoint, the function check_alive() is called, which
    returns a JSON response with the message "ALIVE".
    """
    return jsonify("ALIVE")


def flask_debug():
    """
    The `flask_debug` function starts a Flask application in debug mode without using the reloader.
    Do not use for production for fucks sake.
    """
    logger = logging.getLogger("discord")

    logger.info("Attempt debug")
    # trunk-ignore(bandit/B201)
    app.run(debug=True, use_reloader=False, host=HOST_ADDRESS, port=PORT_ADDRESS)


def waitress_run():
    """
    The `waitress_run` function serves the `app` on host "0.0.0.0"
    and port 5000 using the Waitress WSGI server.
    """
    logger = logging.getLogger("discord")

    logger.info("Attempt waitress")
    serve(app, host=HOST_ADDRESS, port=PORT_ADDRESS)


def scan_queue(stop_event: Optional[threading.Event] = None):
    """
    The function `scan_queue` reads data from a queue, logs it, and appends it to another queue.

    A bounded ``get(timeout=1)`` is used instead of a blocking ``get()`` so the
    worker can observe ``stop_event`` and exit cleanly during shutdown.

    :param stop_event: optional :class:`threading.Event`; when set the loop
    stops at the next iteration.
    """
    logger = logging.getLogger("discord")
    while True:
        if stop_event and stop_event.is_set():
            logger.info("scan_queue: stop requested")
            break
        try:
            data = OUT_COMM_Q.get(timeout=1)
        except Empty:
            continue
        logger.info(data)
        awaiting_q.append(data)


def scan_incoming(stop_event: Optional[threading.Event] = None):
    """
    The `scan_incoming` function continuously checks for incoming data, processes it, and logs when data
    is found.

    :param stop_event: optional :class:`threading.Event`; when set the loop
    stops at the next iteration.
    """
    logger = logging.getLogger("discord")
    while True:
        if stop_event and stop_event.is_set():
            logger.info("scan_incoming: stop requested")
            break
        # Sweep stale health-check pings first: if a librarian is dead the pong
        # never arrives, so drop ping records past their TTL. scan_incoming is
        # the sole remover of awaiting_q, so this needs no lock (scan_queue only
        # appends). Snapshot with list() so removal during iteration is safe.
        now = time.monotonic()
        for record in list(awaiting_q):
            if (
                getattr(record, "is_ping", False)
                and getattr(record, "answered", None) is not None
                and not record.answered.is_set()
                and now - getattr(record, "created", now) > PING_TTL_SECONDS
            ):
                awaiting_q.remove(record)
        try:
            answer = incoming_q.get(block=False)
            logger.info("DATA FOUND")
            # Health-check pong (shape: {"__pong__": uuid}). Close the waiter's
            # event and drop its record. It must NEVER fall through to the
            # result/orphan path below, or the librarian cog would later pull it
            # off IN_COMM_Q and post a bogus "no results" message to Discord.
            if isinstance(answer, dict) and "__pong__" in answer:
                pong_uuid = answer["__pong__"]
                for record in list(awaiting_q):
                    if getattr(record, "uuid", None) == pong_uuid:
                        event = getattr(record, "answered", None)
                        if event is not None:
                            event.set()
                        awaiting_q.remove(record)
                logger.info("PONG matched for %s", pong_uuid)
                continue
            # Collect matched records and drop them from awaiting_q afterwards -
            # they used to stay forever (awaiting_q only ever grew), leaking
            # memory over the bot's uptime and letting a reused UUID re-match a
            # stale record.
            matched = []
            for record in awaiting_q:
                if record.uuid in answer.keys():
                    record.stop = True
                    record.entries = answer[record.uuid]
                    IN_COMM_Q.put(record)
                    matched.append(record)
            for record in matched:
                awaiting_q.remove(record)
            if not matched:
                for key in answer.keys():
                    record = QueryControl("Orphaned", key, "Orphan", None)
                    record.stop = True
                    record.entries = answer[record.uuid]
                    IN_COMM_Q.put(record)
        except Empty:
            time.sleep(1)


def librarian_ping(address: str, endpoint: str, headers: Optional[dict] = None,
                   timeout: float = 3.0) -> bool:
    """Health-check the librarian by round-tripping a ping through the FULL path.

    The ping is a pseudo-query that exercises exactly the same machinery a real
    search does, on BOTH sides:

    * bot out: a ``QueryControl`` rides ``OUT_COMM_Q`` -> ``scan_queue`` ->
      ``awaiting_q`` just like a real query,
    * librarian: it is POSTed to the librarian, which must pull it off its OWN
      internal queue and answer WITHOUT running a search,
    * bot in: the pong comes back over ``/conjurer`` -> ``incoming_q`` ->
      ``scan_incoming``, which matches it by uuid and sets our event.

    Returns True only when that whole loop closes within ``timeout``. Never
    blocks longer than roughly ``timeout`` and cannot deadlock: the POST is
    bounded, the wait is bounded, and a ping whose pong never arrives is swept
    out of ``awaiting_q`` by ``scan_incoming`` (PING_TTL_SECONDS).
    """
    logger = logging.getLogger("discord")
    ping_uuid = str(uuidlib.uuid4())
    answered = threading.Event()
    query = QueryControl("healthcheck", ping_uuid, "__ping__", None)
    query.is_ping = True
    query.answered = answered
    query.created = time.monotonic()
    # Enter the bot-side comm queue BEFORE the POST, so the record is already in
    # awaiting_q by the time the pong can come back (no lost-wakeup race).
    OUT_COMM_Q.put(query)
    try:
        response = requests.post(
            f"{address}{endpoint}",
            json={"UUID": ping_uuid},
            headers=headers or {},
            timeout=timeout,
        )
    except requests.exceptions.RequestException as exc:
        logger.info("Librarian ping POST failed (%s): %s", ping_uuid, exc)
        return False  # stale record is swept by scan_incoming
    if response.status_code != 200:
        logger.info("Librarian ping rejected (%s): HTTP %s", ping_uuid, response.status_code)
        return False
    if answered.wait(timeout):
        logger.info("Librarian ping round-trip OK (%s)", ping_uuid)
        return True
    logger.info("Librarian ping timed out after %ss (%s)", timeout, ping_uuid)
    return False


def get_stream_title(tag: bytes) -> str:
    title = ""
    if m := SRCHTITLE(tag):
        # decode, strip, unescape and remove surrounding quotes (may not even be the same type of quote)
        title = m.group("title").decode("utf-8").strip().replace("\\", "")[1:-1]
    return title


def id3(url: str) -> dict:
    request = urequest.Request(url, headers={"Icy-MetaData": 1})

    with urequest.urlopen(request) as resp:
        metaint = int(resp.headers.get("icy-metaint", "-1"))
        if metaint < 0:
            return False
        resp.read(
            metaint
        )  # this isn't seekable so, arbitrarily read to the point we want
        # Guard the headers: an Icecast stream that omits icy-name / icy-genre
        # (e.g. while the radio is down) made `.title()` raise AttributeError on
        # None, 500-ing the /prepped_tracks "next" handler that calls this.
        tagdata = dict(
            site_url=resp.headers.get("icy-url"),
            name=(resp.headers.get("icy-name") or "").title(),
            genre=(resp.headers.get("icy-genre") or "").title(),
            title=get_stream_title(resp.read(255)),
        )
        return tagdata


def comm_subroutine(stop_event: Optional[threading.Event] = None):
    """
    The `comm_subroutine` function starts multiple threads to run different tasks concurrently.

    Workers run as daemon threads and honour an optional ``stop_event`` so the
    bot can shut the communication layer down cleanly instead of blocking
    forever on ``join()``.

    :param stop_event: optional :class:`threading.Event` shared with the caller
    to coordinate a cooperative shutdown.
    """
    logger = logging.getLogger("discord")
    logger.setLevel(logging.DEBUG)
    logger.info("Started comms")
    threads = []
    # NOTE: flask_debug is the dev server bound to the SAME host:port as
    # waitress - running both kills the comm layer with 'address in use'.
    # Enable it only INSTEAD of waitress_run, never alongside.
    # threads.append(threading.Thread(target=flask_debug))
    threads.append(threading.Thread(target=waitress_run, daemon=True))
    threads.append(
        threading.Thread(
            target=scan_queue, kwargs={"stop_event": stop_event}, daemon=True
        )
    )
    threads.append(
        threading.Thread(
            target=scan_incoming, kwargs={"stop_event": stop_event}, daemon=True
        )
    )

    for worker in threads:
        worker.start()

    try:
        while any(thread.is_alive() for thread in threads):
            if stop_event and stop_event.is_set():
                break
            time.sleep(0.5)
    finally:
        if stop_event:
            stop_event.set()


if __name__ == "__main__":
    comm_subroutine()