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

# -*- coding: utf-8 -*-
"""SVG -> PNG 转换 + 可选 HD2 风格边框，用于战备图标上传。"""

from __future__ import annotations

import base64
import io
import re
from pathlib import Path

from PIL import Image, ImageDraw

SCHEMA_SIZE = 550  # 与随包 PNG 图标一致
TYPE_COLORS = {
    "blue": "#6bb7ea",
    "red": "#fe6d6a",
    "green": "#7ecb3f",
}

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"


def _hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
    h = hex_color.lstrip("#")
    if len(h) == 3:
        h = "".join(c * 2 for c in h)
    return tuple(int(h[i : i + 2], 16) for i in (0, 2, 4))  # type: ignore[return-value]


def _apply_frame(src_png: bytes, size: int, frame_color: str) -> bytes:
    """把透明的 SVG 画面叠到 HD2 风格画布上：彩色边框 + 深色内底 + 居中图形。"""
    rgb = _hex_to_rgb(frame_color)
    img = Image.open(io.BytesIO(src_png)).convert("RGBA")
    canvas = Image.new("RGBA", (size, size), rgb + (255,))
    draw = ImageDraw.Draw(canvas)
    border = max(10, size // 20)
    # 外圈用类型色，内部深色
    draw.rectangle([0, 0, size - 1, size - 1], fill=rgb + (255,))
    draw.rectangle([border, border, size - 1 - border, size - 1 - border], fill=(12, 12, 12, 255))
    # 去掉透明边距再居中
    alpha = img.getchannel("A")
    bbox = alpha.getbbox()
    glyph = img.crop(bbox) if bbox else img
    inner = size - 2 * border
    avail = int(inner * 0.82)
    glyph.thumbnail((avail, avail), Image.LANCZOS)
    ox = (size - glyph.width) // 2
    oy = (size - glyph.height) // 2
    canvas.alpha_composite(glyph, (ox, oy))
    out = io.BytesIO()
    canvas.save(out, format="PNG")
    return out.getvalue()


async def _render_svg(svg_text: str, size: int) -> bytes:
    """用 Chromium 的 canvas 把 SVG 光栅化为带透明通道的 PNG。

    用 toDataURL 导出而非截图，避免 omit_background 在部分环境把透明背景
    拍成白色，从而保住透明 alpha。
    """
    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()
    return base64.b64decode(data_url.split(",", 1)[1])


async def svg_to_png(svg_text: str, framed: bool = False, frame_color: str | None = None, size: int = SCHEMA_SIZE) -> bytes:
    """SVG -> PNG。framed=True 时叠加 HD2 风格边框到类型色。"""
    if not svg_text.strip():
        raise ValueError("SVG 内容为空")
    png = await _render_svg(svg_text, size)
    if framed:
        png = _apply_frame(png, size, frame_color or TYPE_COLORS["blue"])
    return png


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}"