HeXi/hexi/plugins/nonebot_plugin_deer_pipe/utils/render.py

"""打卡记鹿 —— 渲染工具(utils 层)。"""

from __future__ import annotations

import calendar
import os
import random
from collections import Counter
from datetime import datetime, timedelta

from nonebot.log import logger
from PIL import Image, ImageDraw, ImageFont, ImageOps

from ..repository import get_records

# 资源统一放插件 res/ 下
filepath = os.path.join(os.path.dirname(os.path.dirname(__file__)), "res").replace("\\", "/")


async def draw_calendar_grid_image(uid, year, month):
    first_weekday, num_days = calendar.monthrange(year, month)
    first_weekday = (first_weekday + 1) % 7
    cell_size = 50
    cell_padding = 5
    num_rows = (num_days + first_weekday + 6) // 7 + 1
    image_width = 7 * (cell_size + cell_padding) - cell_padding
    image_height = num_rows * (cell_size + cell_padding) - cell_padding
    img = Image.new("RGBA", (image_width, image_height), "white")
    draw = ImageDraw.Draw(img)
    font_path = filepath + "/font/SourceHanSansCN-Medium.otf"
    logger.debug(f"打卡字体路径: {font_path}")
    font_normal = ImageFont.truetype(font=font_path, size=18)
    font_large = ImageFont.truetype(font=font_path, size=24)
    for i in range(7):
        x = i * (cell_size + cell_padding)
        draw.line([(x, 55), (x, image_height)], fill="black", width=1)
    for j in range(0, image_height + cell_size, cell_size + cell_padding):
        draw.line([(0, j), (image_width, j)], fill="black", width=1)
    deer_pipe_path = filepath + "/img/deer_pipe.jpg"
    deer_pipe_img = Image.open(deer_pipe_path)
    deer_pipe_img = deer_pipe_img.resize((54, 50))
    title = f"{month}月打卡记鹿"
    text_length = draw.textlength(text=title, font=font_large)
    img_xy = (int(((image_width - text_length + 55) / 2) - 55), 5)
    title_pos = ((image_width - text_length + 55) / 2, 10)
    img.paste(deer_pipe_img, img_xy)
    draw.text(title_pos, title, fill="black", font=font_large)
    day_counts = Counter(await get_records(uid))
    date = datetime(year, month, 1)
    for day in range(1, num_days + 1):
        weekday = (first_weekday + day - 1) % 7
        row = (first_weekday + day - 1) // 7 + 1
        x = weekday * (cell_size + cell_padding) + cell_padding
        y = row * (cell_size + cell_padding) + cell_padding
        img.paste(deer_pipe_img, (x - 4, y))
        draw.text((x + 20, y - 2), f"{str(day)}", fill="black", font=font_normal)
        for i in range(day_counts.get(day, 0)):
            im = get_random_right()
            img = image_paste(im, img, (x - 4 + i * 10, y))
    img = ImageOps.expand(img, border=5, fill="black")
    img = ImageOps.expand(img, border=5, fill="white")
    return img


selected_images = []


def get_random_right():
    global selected_images
    path = filepath + "/img/right"
    im_name = [name for name in os.listdir(path) if name not in selected_images]
    if not im_name:
        selected_images = []
        im_name = os.listdir(path)
    index = random.randint(0, len(im_name) - 1)
    im = im_name[index]
    selected_images.append(im)
    im_path = os.path.join(path, im)
    im_file = Image.open(im_path)
    im_file = im_file.resize((50, 50))
    return im_file


def image_paste(paste_image, under_image, pos):
    r, g, b, alpha = paste_image.split()
    under_image.paste(paste_image, pos, alpha)
    return under_image