HeXi/hexi/plugins/nonebot_plugin_makeaquote/make_a_qoute.py

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 .Reply import Reply

# --- 常量 ---
_FONT_DIR = Path(__file__).resolve().parent / "data" / "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]:
    """
    自动缩放字号使文本适合 max_width × max_height 区域。
    返回 (行列表, 字体, 字号, 行间距)。
    """
    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          # ≈ 640px
    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)

    # --- Pilmoji ---
    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]:
    """
    按宽度自动换行，支持中日韩文本。
    先按换行符分行，再用 cjk_textwrap.wrap 按宽度折行。
    """
    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):
    """
    为图像添加从 x_s 到 x_e 的左右渐变透明度遮罩。
    使用 numpy 加速，避免逐像素 putpixel 的低效操作。
    """
    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:
    """
    下载指定 URL 内容，带 3 次重试和超时。
    如果失败，返回空字节串 b""。
    """
    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""