republisher/repub/web.py

from __future__ import annotations

import asyncio
import hashlib
from collections.abc import AsyncGenerator
from contextlib import suppress
from typing import cast

import htpy as h
from datastar_py.quart import DatastarResponse
from datastar_py.sse import DatastarEvent
from htpy import Renderable
from quart import Quart, Response, request, url_for

from repub.datastar import RefreshBroker, render_stream
from repub.pages import admin_component, shim_page

REFRESH_BROKER_KEY = "repub.refresh_broker"
ACTIVE_JOBS_KEY = "repub.demo_active_jobs"
REFRESH_TASK_KEY = "repub.demo_refresh_task"


def _render_shim_page(*, stylesheet_href: str, datastar_src: str) -> tuple[str, str]:
    head = (
        h.title["Republisher Admin UI"],
        h.link(rel="stylesheet", href=stylesheet_href),
    )
    body = str(shim_page(datastar_src=datastar_src, head=head))
    etag = hashlib.sha256(body.encode("utf-8")).hexdigest()
    return body, etag


def create_app(*, enable_demo_refresh: bool = True) -> Quart:
    app = Quart(__name__)
    app.extensions[REFRESH_BROKER_KEY] = RefreshBroker()
    app.extensions[ACTIVE_JOBS_KEY] = 12

    if enable_demo_refresh:

        @app.before_serving
        async def start_demo_refresh() -> None:
            app.extensions[REFRESH_TASK_KEY] = asyncio.create_task(
                _demo_refresh_loop(app)
            )

        @app.after_serving
        async def stop_demo_refresh() -> None:
            task = cast(asyncio.Task[None] | None, app.extensions.get(REFRESH_TASK_KEY))
            if task is None:
                return
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task

    @app.get("/")
    async def index() -> Response:
        body, etag = _render_shim_page(
            stylesheet_href=url_for("static", filename="app.css"),
            datastar_src=url_for("static", filename="datastar@1.0.0-RC.8.js"),
        )
        if request.if_none_match.contains(etag):
            response = Response(status=304)
            response.set_etag(etag)
            return response

        response = Response(body, mimetype="text/html")
        response.set_etag(etag)
        return response

    @app.post("/")
    async def index_patch() -> DatastarResponse:
        queue = get_refresh_broker(app).subscribe()
        stream = render_stream(
            queue,
            render=lambda: render_dashboard(app),
            last_event_id=request.headers.get("last-event-id"),
        )
        return DatastarResponse(_unsubscribe_on_close(queue, stream, app))

    return app


def get_refresh_broker(app: Quart) -> RefreshBroker:
    return cast(RefreshBroker, app.extensions[REFRESH_BROKER_KEY])


def trigger_refresh(app: Quart, event: object = "refresh-event") -> None:
    get_refresh_broker(app).publish(event)


async def render_dashboard(app: Quart) -> Renderable:
    return admin_component(active_jobs=str(get_active_jobs(app)))


async def _unsubscribe_on_close(
    queue: object, stream: AsyncGenerator[DatastarEvent, None], app: Quart
) -> AsyncGenerator[DatastarEvent, None]:
    try:
        async for event in stream:
            yield event
    finally:
        get_refresh_broker(app).unsubscribe(cast(asyncio.Queue[object], queue))


def get_active_jobs(app: Quart) -> int:
    return cast(int, app.extensions[ACTIVE_JOBS_KEY])


def set_active_jobs(app: Quart, value: int) -> None:
    app.extensions[ACTIVE_JOBS_KEY] = value


async def _demo_refresh_loop(app: Quart) -> None:
    while True:
        await asyncio.sleep(1)
        set_active_jobs(app, get_active_jobs(app) + 1)
        trigger_refresh(app)