app/api/tex.py

"""LaTeX 公式图片下发 / 按需抽取端点。

手表端 ``<image>`` + ``request.download`` 都无法携带自定义 header（平台限制），
所以沿用生图 / TTS 的 capability URL 思路：URL 里带内容寻址 id，
路径里的 id = sha1(latex|fs|fg) 前 20 位，服务端校验 id 与内容一致才渲染。

三个端点：
- ``GET  /v1/xtc/tex/{id}.png``   取 PNG（内容寻址，字节永不变化 → 可长缓存）
- ``POST /v1/xtc/tex/prepare``    批量预热：端上把「本地缓存缺失」的公式一次性送过来，
                                  服务端批量渲染填进进程内 LRU，之后逐个下载短 URL 即可。
                                  这样 URL 里不需要塞 LaTeX（避免超长 URL），
                                  也是「先查端上缓存，没有再请求后端」的服务端一半。
- ``POST /v1/xtc/tex/extract``    按需抽取：历史消息里没有公式清单时（功能上线前的老消息、
                                  或后端当时渲染失败），端上把「当前聚焦那条消息」的正文
                                  送过来，服务端重新抽取并渲染，回一份 texs 让端上落盘缓存。

为什么 ``{id}.png`` 要支持 ``?f=``：
  进程内 LRU 在 Space 重建 / 重启后必然为空。若 URL 里不带公式内容，重启后老消息的
  公式就永久 404（这正是改造前「很多公式渲染不出来」的根因之一）。
  带上 ``?f=``（以及 fs/fg）后，LRU 未命中也能按内容校验并重新渲染，URL 永久自洽。
"""
from __future__ import annotations

import asyncio
import logging
import time

from fastapi import APIRouter, Depends, Query, Request, Response

from ..auth import require_access_key
from ..errors import HttpError
from ..services import formula
from ._common import CORS_HEADERS, ok_with_cors, read_json_body

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/v1/xtc/tex", tags=["formula"])

# 每分钟最多按需渲染多少次（命中 LRU 不计入）。防公开端点被刷 CPU。
_RENDER_BUDGET = 600
_window_start = 0.0
_window_used = 0


def _allow_render() -> bool:
    global _window_start, _window_used
    now = time.time()
    if now - _window_start >= 60.0:
        _window_start = now
        _window_used = 0
    if _window_used >= _RENDER_BUDGET:
        return False
    _window_used += 1
    return True


@router.get("/{tex_id}.png")
async def get_tex_png(
    tex_id: str,
    f: str = Query(default="", description="LaTeX 正文（不带 $）"),
    fs: int = Query(default=formula.DEFAULT_FS),
    fg: str = Query(default=formula.DEFAULT_FG),
) -> Response:
    """返回公式 PNG（透明底）。端上用 request.download 拉取后落盘缓存。"""
    entry = formula.get_png(tex_id)
    if entry is None:
        # LRU 未命中：校验内容寻址 id 后按需渲染（Space 重启后老消息仍可显示）
        latex = str(f or "").strip()
        if not latex or len(latex) > formula._MAX_LATEX_LEN:
            raise HttpError("tex not found", status=404, code="not_found")
        if formula.formula_id(latex, fs, fg) != tex_id:
            raise HttpError("tex id mismatch", status=403, code="forbidden")
        if not _allow_render():
            raise HttpError("render rate limited", status=429, code="rate_limited")
        png = await asyncio.to_thread(formula.render_png, latex, fs, formula.DEFAULT_SCALE, fg)
        if png is None:
            raise HttpError("tex render failed", status=404, code="not_found")
        entry = png

    return Response(
        content=entry,
        media_type="image/png",
        headers={
            **CORS_HEADERS,
            # 内容寻址 → 同一 id 的字节永不变化，可长缓存
            "Cache-Control": "public, max-age=31536000, immutable",
        },
    )


@router.post("/prepare")
async def prepare_texs(request: Request, _key: str = Depends(require_access_key)):
    """批量预热公式 PNG 到进程内缓存。

    body: ``{"items": [{"id": "<id>", "latex": "...", "fs": 12}]}``
    返回: ``{"ok": true, "ready": [<id>...], "failed": [{"id","error"}...], "texs": [...]}``

    端上流程：本地 ``file.access`` 命中就跳过；未命中的 id 一次性送这里 → 成功后用
    ``{id}.png`` 短 URL 并发下载。
    """
    body = await read_json_body(request)
    raw = body.get("items")
    items = raw if isinstance(raw, list) else []
    if not items:
        return ok_with_cors({"ready": [], "failed": [], "texs": []})
    if len(items) > formula._MAX_FORMULAS * 2:
        items = items[: formula._MAX_FORMULAS * 2]

    fs_default = int(body.get("fs") or formula.DEFAULT_FS)
    wanted: list[dict] = []
    bad: list[dict] = []
    for it in items:
        if not isinstance(it, dict):
            continue
        latex = str(it.get("latex") or "").strip()
        tid = str(it.get("id") or "").strip()
        fs = int(it.get("fs") or fs_default)
        if not latex or len(latex) > formula._MAX_LATEX_LEN:
            bad.append({"id": tid, "error": "latex 为空或过长"})
            continue
        expect = formula.formula_id(latex, fs, formula.DEFAULT_FG)
        if tid and tid != expect:
            bad.append({"id": tid, "error": "id 与内容不匹配"})
            continue
        wanted.append({"latex": latex, "fs": fs, "id": expect})

    if not formula.is_available():
        raise HttpError("公式渲染不可用", status=503, code="service_unavailable")

    infos = await asyncio.to_thread(
        formula.render_many, [{"latex": w["latex"], "fs": w["fs"]} for w in wanted]
    )

    ready: list[str] = []
    failed: list[dict] = list(bad)
    texs: list[dict] = []
    for w in wanted:
        info = infos.get(w["id"])
        if not info:
            failed.append({"id": w["id"], "error": "渲染失败"})
            continue
        ready.append(w["id"])
        # 返回 w/h，端上不必再单独问尺寸。url 用短路径：
        # 端上在下载前会先调本端点把 LRU 预热，所以短 URL 足够；
        # 用短 URL 也能把会话 JSON 里每条公式的体积压到 ~40B（手表存储紧张）。
        texs.append({
            "id": w["id"],
            "latex": w["latex"],
            "w": info["w"],
            "h": info["h"],
            "url": "/v1/xtc/tex/" + w["id"] + ".png",
        })
    return ok_with_cors({"ready": ready, "failed": failed, "texs": texs})


@router.post("/extract")
async def extract_texs(request: Request, _key: str = Depends(require_access_key)):
    """按需抽取 + 渲染正文里的公式（历史消息没有公式清单时用）。

    body: ``{"text": "<消息正文>", "fs": 12, "limit": 16}``
    返回: ``{"ok": true, "texs": [{id, latex, src, start, end, w, h, url}], "count": n}``

    设计意图：端上只把「当前正在看的那条消息」送过来，避免为整段历史批量渲染
    （手表下载管理器有 70 个文件的限制，服务端渲染也是纯 CPU 开销）。
    """
    body = await read_json_body(request)
    text = str(body.get("text") or "")
    if not text.strip():
        return ok_with_cors({"texs": [], "count": 0})
    # 正文长度上限：单条消息最多 ~20000 字，超出的截断（抽取/渲染只针对可见内容）
    if len(text) > 20000:
        text = text[:20000]
    fs = int(body.get("fs") or formula.DEFAULT_FS)
    try:
        limit = int(body.get("limit") or formula._MAX_FORMULAS)
    except Exception:
        limit = formula._MAX_FORMULAS
    limit = max(1, min(formula._MAX_FORMULAS, limit))

    if not formula.is_available():
        raise HttpError("公式渲染不可用", status=503, code="service_unavailable")

    texs = await asyncio.to_thread(formula.build_texs, text, fs, limit)
    return ok_with_cors({"texs": texs, "count": len(texs)})


@router.get("/stats")
async def tex_stats(_key: str = Depends(require_access_key)):
    """诊断：侧车是否就绪、缓存条数、上限配置。"""
    return ok_with_cors({"stats": formula.stats()})