HeXi/hexi/plugins/nonebot_plugin_helldivers_tools/utils/icon_utils.py

# -*- coding: utf-8 -*-
"""SVG -> PNG 转换。

与 docs/HD2/run.py 一致：把上传的单张 SVG 按输出尺寸（默认 550x550）栅格化成
透明背景 PNG，不做批量，也不叠加边框。

优先用 cairosvg（与 run.py 完全一致）；若运行环境缺 cairo 原生库，则回退到
Chromium canvas 渲染（透明背景、同尺寸），保证 Web 后台上传始终可用。
"""

from __future__ import annotations

import base64
import logging
import re
from pathlib import Path

logger = logging.getLogger(__name__)

try:
    import cairosvg  # type: ignore
    _HAS_CAIROSVG = True
except Exception:  # noqa: BLE001
    cairosvg = None
    _HAS_CAIROSVG = False

SCHEMA_SIZE = 550  # 与随包 PNG 图标一致

basic_path = Path(__file__).resolve().parents[1]
IMG_DIR = basic_path / "res" / "img" / "helldivers"


def sanitize_filename(name: str) -> str:
    """把战备名转成安全的 Windows 文件名（保留中文/空格，去掉非法字符）。"""
    s = re.sub(r'[\/:*?"<>|]', "_", name).strip()
    s = re.sub(r"\s+", " ", s)
    return s or "strategem"


async def _render_svg_canvas(svg_text: str, size: int) -> bytes:
    """用 Chromium canvas 把 SVG 光栅化为带透明通道的 PNG（兜底方案）。"""
    from playwright.async_api import async_playwright

    data_uri = "data:image/svg+xml;base64," + base64.b64encode(svg_text.encode("utf-8")).decode()
    script = (
        "() => new Promise((resolve) => {"
        "const img = new Image();"
        "img.onload = () => {"
        "const canvas = document.createElement('canvas');"
        f"canvas.width = {size}; canvas.height = {size};"
        "const ctx = canvas.getContext('2d');"
        f"ctx.clearRect(0,0,{size},{size});"
        f"const scale = Math.min({size} / img.width, {size} / img.height);"
        "const w = img.width * scale, h = img.height * scale;"
        f"ctx.drawImage(img, ({size}-w)/2, ({size}-h)/2, w, h);"
        "resolve(canvas.toDataURL('image/png'));"
        "};"
        "img.onerror = () => resolve('');"
        f"img.src = '{data_uri}';"
        "})"
    )
    async with async_playwright() as p:
        browser = await p.chromium.launch()
        page = await browser.new_page()
        await page.set_content("<!-- svg to png -->")
        data_url = await page.evaluate(script)
        await browser.close()
    if not data_url:
        raise ValueError("SVG 转换失败")
    return base64.b64decode(data_url.split(",", 1)[1])


async def svg_to_png(svg_text: str, size: int = SCHEMA_SIZE) -> bytes:
    """SVG -> PNG（透明背景，默认 550x550）。

    参考 docs/HD2/run.py：cairosvg.svg2png 按 output_width/output_height 栅格化到目标尺寸。
    """
    if not svg_text.strip():
        raise ValueError("SVG 内容为空")
    if _HAS_CAIROSVG:
        try:
            return cairosvg.svg2png(
                bytestring=svg_text.encode("utf-8"),
                output_width=size,
                output_height=size,
            )
        except Exception as e:  # noqa: BLE001
            logger.warning(f"cairosvg 转换失败，改用 canvas 兜底: {type(e).__name__}: {e}")
    return await _render_svg_canvas(svg_text, size)


def save_icon_png(png_bytes: bytes, name: str) -> str:
    """写入 img/helldivers/<sanitized name>.png，返回相对路径。"""
    IMG_DIR.mkdir(parents=True, exist_ok=True)
    filename = sanitize_filename(name) + ".png"
    dest = IMG_DIR / filename
    dest.write_bytes(png_bytes)
    return f"res/img/helldivers/{filename}"