conjurer/communication_subroutine.py

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

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()
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
        try:
            answer = incoming_q.get(block=False)
            logger.info("DATA FOUND")
            # 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 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()