lair.moe/blueprints/risdeveau/modules/api/lb.py

import atexit
import re
from dataclasses import dataclass
from hashlib import md5
from json import dumps
from time import time
from urllib.parse import parse_qs, urlparse

from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.triggers.interval import IntervalTrigger
import requests
from flask import Flask, jsonify

from .scheduler_guard import should_start_scheduler
from .shared_cache import atomic_write_json, cache_file, load_json_if_newer


@dataclass
class Cache:
    data = {}
    last_updated = time()
    status = None


data = {
    "caches": {
        "now": Cache(),
        "listens": Cache()
    },
    "last_updated": time(),
    "etag": ""
}

_IS_WRITER = should_start_scheduler()
_CACHE_PATH = cache_file("listenbrainz")
_CACHE_MTIME = 0.0


def refresh_cache() -> None:
    """Refresh in-memory cache from the shared JSON file (for non-writer workers)."""
    global _CACHE_MTIME
    if _IS_WRITER:
        return

    payload, mtime = load_json_if_newer(_CACHE_PATH, _CACHE_MTIME)
    if payload is None:
        return

    try:
        data["etag"] = payload.get("etag", data.get("etag", ""))
        data["last_updated"] = payload.get("last_updated", data.get("last_updated", time()))
        caches = payload.get("caches", {})
        for key, cache in data.get("caches", {}).items():
            if isinstance(caches, dict) and key in caches and isinstance(caches[key], dict):
                c = caches[key]
                cache.data = c.get("data", cache.data)
                cache.last_updated = c.get("last_updated", cache.last_updated)
                cache.status = c.get("status", cache.status)
    finally:
        _CACHE_MTIME = mtime


def _persist_cache() -> None:
    """Persist current cache state to the shared JSON file (writer worker only)."""
    if not _IS_WRITER:
        return

    payload = {
        "etag": data.get("etag", ""),
        "last_updated": data.get("last_updated", time()),
        "caches": {
            k: {
                "data": v.data,
                "last_updated": v.last_updated,
                "status": v.status,
            }
            for k, v in data.get("caches", {}).items()
        },
    }
    atomic_write_json(_CACHE_PATH, payload)


def yt_cover(youtube_url):
    parsed_url = urlparse(youtube_url)

    if parsed_url.netloc in ("youtube.com", "music.youtube.com"):
        query_params = parse_qs(parsed_url.query)
        video_id = query_params.get('v', [None])[0]

    elif parsed_url.netloc == 'youtu.be':
        video_id = parsed_url.path[1:]

    if not video_id:
        return

    return f"https://img.youtube.com/vi/{video_id}/2.jpg"


def parse_listens(json: dict) -> dict:
    cover_replacing = {
        "1e699948-c7c8-4bb2-9f8b-62e14b882a5d": "ca464c1d-5848-45bb-b92d-b1e4b00f9d65",
        "0d516a93-061e-4a27-9cf7-f36e3a96f888": "5cc0c0c7-22f9-4a4b-a24c-f1a6732f813b",
        "92ea5cc8-80e0-4da0-a10b-1bc2f8e8781e": "e8f3e14a-4794-4bab-b403-d562cdad4c2f",
    }

    new_json = {
        "count": json["count"],
        "listens": []
    }

    for track in json["listens"]:
        listened_at = track.get("listened_at", 0)
        track = track["track_metadata"]

        new_track = {
            "artist_name": track["artist_name"],
            "track_name": track["track_name"],
            "listened_at": listened_at
        }

        if mb := track.get("mbid_mapping"):
            new_track["id"] = \
                cover_replacing.get(mb["release_mbid"],
                mb.get("caa_release_mbid",
                mb["release_mbid"]
            ))
            new_track["artist_name"] = mb["artists"][0]["artist_credit_name"]
            new_track["track_name"] = mb["recording_name"]
        elif info := track.get("additional_info"):
            if info \
                .get("music_service_name", "") \
                .lower() in ("youtube", "youtube music"):
                    if cover := yt_cover(track["additional_info"]["origin_url"]):
                        new_track["cover_url"] = cover

        if "cover_url" not in new_track.keys() and "id" in new_track.keys():
            new_track["cover_url"] = "/asset/mb/" + new_track["id"]

        new_json["listens"].append(new_track)

    return new_json


def api_request(url: str, cache: Cache):
    changed = False
    prev_status = cache.status

    try:
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            json = parse_listens(response.json().get("payload"))
            cache.status = 'success'

            if cache.data != json:
                cache.data = json
                cache.last_updated = time()
                data['last_updated'] = time()
                data['etag'] = md5(''.join(
                    ( dumps(data['caches'][x].data) for x in data['caches'] )
                ).encode()).hexdigest()
                changed = True
        else:
            cache.status = f'error: {response.status_code}'
    except Exception as e:
        cache.status = f'error: {str(e)}'

    if prev_status != cache.status:
        changed = True

    if changed:
        _persist_cache()


scheduler = BackgroundScheduler()

if _IS_WRITER:
    scheduler.add_job(
        func=lambda: api_request("https://api.listenbrainz.org/1/user/risdeveau/listens?count=5", data['caches']['listens']),
        trigger=IntervalTrigger(minutes=1),
        id='risdeveau.listenbrainz.listens',
        replace_existing=True
    )
    scheduler.add_job(
        func=lambda: api_request("https://api.listenbrainz.org/1/user/risdeveau/playing-now", data['caches']['now']),
        trigger=IntervalTrigger(seconds=15),
        id='risdeveau.listenbrainz.playing-now',
        replace_existing=True
    )
    _persist_cache()
    scheduler.start()
    atexit.register(lambda: scheduler.shutdown())
else:
    refresh_cache()