HeXi/hexi/plugins/nonebot_plugin_makeaquote/services/generate.py

"""语录 —— 图片生成(Service 层)。"""

from __future__ import annotations

import hashlib
from io import BytesIO
from pathlib import Path
from time import localtime, strftime

import cjk_textwrap
import httpx
import numpy as np
from nonebot.log import logger
from PIL import Image, ImageDraw, ImageFont
from pilmoji import Pilmoji
from pilmoji.source import TwemojiEmojiSource

from ..utils.reply import Reply

# 字体资源在插件 res/font 下
_FONT_DIR = Path(__file__).resolve().parents[1] / "res" / "font"
FONT_SMALL = 24
FONT_TIME = 18
FONT_TEXT_MAX = 32
FONT_TEXT_MIN = 12

_font_cache: dict = {}


def _get_font(name: str, size: int) -> ImageFont.FreeTypeFont:
    key = (name, size)
    if key not in _font_cache:
        path = _FONT_DIR / name
        _font_cache[key] = ImageFont.truetype(str(path), size, encoding="unic")
    return _font_cache[key]


def _font_small() -> ImageFont.FreeTypeFont:
    return _get_font("SourceHanSansCN-Bold.otf", FONT_SMALL)


def _font_time() -> ImageFont.FreeTypeFont:
    return _get_font("SourceHanSansCN-Bold.otf", FONT_TIME)


def _auto_fit(
    text: str,
    draw: ImageDraw.ImageDraw,
    max_width: float,
    max_height: float,
) -> tuple[list[str], ImageFont.FreeTypeFont, int, int]:
    for size in range(FONT_TEXT_MAX, FONT_TEXT_MIN - 1, -2):
        font = _get_font("SourceHanSansCN-Medium.otf", size)
        chars_per_line = max(1, int(max_width / (size * 0.92)))
        lines = wrap_text(text, chars_per_line)
        spacing = int(size * 0.75)
        total_h = len(lines) * (size + spacing) - spacing
        if total_h <= max_height and all(
            draw.textlength(line, font) <= max_width for line in lines
        ):
            return lines, font, size, spacing
    font = _get_font("SourceHanSansCN-Medium.otf", FONT_TEXT_MIN)
    chars_per_line = max(1, int(max_width / (FONT_TEXT_MIN * 0.92)))
    lines = wrap_text(text, chars_per_line)
    spacing = int(FONT_TEXT_MIN * 0.75)
    return lines, font, FONT_TEXT_MIN, spacing


async def generate(reply: Reply) -> BytesIO:
    image = Image.new("RGBA", (1280, 640), (0, 0, 0, 255))
    draw = ImageDraw.Draw(image)
    x_start = 601.1
    x_end = 1241.1
    max_width = x_end - x_start
    max_height = 400
    text_wrapped, text_font, font_size, line_spacing = _auto_fit(
        reply.message, draw, max_width, max_height,
    )
    total_h = len(text_wrapped) * (font_size + line_spacing) - line_spacing
    y_min = 100 - font_size
    y_max = 512
    y = int((y_max - y_min) / 2 - total_h / 2)
    with Pilmoji(image, source=TwemojiEmojiSource, emoji_position_offset=(0, 10)) as emoji:
        for line in text_wrapped:
            text_size = draw.textlength(line, text_font)
            x = int(x_start + (max_width - text_size) / 2)
            emoji.text((x, y), line, font=text_font, fill=(255, 255, 255, 255), embedded_color=True)
            y += font_size + line_spacing
    img_content = await download_avatar(str(reply.user_id))
    avatar = Image.open(img_content).resize((640, 640))
    avatar = alpha_gradient(avatar, 320, 640)
    image.paste(avatar, (0, 0), avatar)
    x_center = 921.1
    card_size = draw.textlength("-  " + reply.user_card, _font_small())
    x_card_start = x_center - card_size / 2
    draw.text((x_card_start, y), "- " + reply.user_card, font=_font_small(), fill=(169, 172, 184, 255))
    fmt_time = strftime("%Y.%m.%d %H:%M", localtime(reply.time))
    time_size = draw.textlength(fmt_time, _font_small())
    x_time_start = x_end - time_size
    draw.text((x_time_start, 640 - FONT_SMALL - 20), fmt_time, font=_font_time(), fill=(169, 172, 184, 255))
    bytes_io = BytesIO()
    image.convert("RGB").save(bytes_io, format="JPEG")
    return bytes_io


def wrap_text(text: str, width: int) -> list[str]:
    return [wrapped for line in text.splitlines() for wrapped in cjk_textwrap.wrap(line, width)]


def alpha_gradient(image, x_s: int, x_e: int):
    image = image.convert("RGBA")
    w, h = image.size
    gradient = np.full((h, w), 255, dtype=np.uint8)
    xs = np.linspace(0, 1, x_e - x_s, endpoint=False)
    alphas = 255 - (np.sin(xs * np.pi - np.pi / 2) / 2 + 0.5) * 255
    alphas = alphas.astype(np.uint8)
    gradient[:, x_s:x_e] = alphas[np.newaxis, :]
    mask = Image.fromarray(gradient, mode="L")
    image.putalpha(mask)
    return image


async def download_avatar(user_id: str) -> BytesIO:
    url = f"http://q1.qlogo.cn/g?b=qq&nk={user_id}&s=640"
    data = await download_url(url)
    if not data or hashlib.md5(data).hexdigest() == "acef72340ac0e914090bd35799f5594e":
        url = f"http://q1.qlogo.cn/g?b=qq&nk={user_id}&s=100"
        data = await download_url(url)
    return BytesIO(data)


async def download_url(url: str) -> bytes:
    async with httpx.AsyncClient(timeout=5.0) as client:
        for i in range(3):
            try:
                resp = await client.get(url)
                if resp.status_code == 200:
                    return resp.content
            except Exception as e:
                logger.warning(f"下载失败 {url}, 重试 {i + 1}/3: {e}")
    return b""