HeXi/hexi/plugins/nonebot_plugin_bf_bot/utils/images.py

"""图像处理工具：缩放/粘贴、头像下载缓存、物品图标、公告栏与标签绘制。

静态资源统一走 paths.RES_DIR；QQ 头像缓存落在 hexi/data/nonebot_plugin_bf_bot/avatar_cache。
"""

import hashlib
import os
from io import BytesIO

import cv2
import numpy as np
import requests
import requests.exceptions
from PIL import Image
from nonebot import logger

from ..paths import AVATAR_CACHE_DIR, RES_DIR
from ..repository import DesignationManager, UserManager
from .params import rank_pic


def png_resize(source_file, new_width=0, new_height=0, resample="LANCZOS", ref_file=""):
    """
    PNG 等比缩放（保持 Alpha 透明度）
    :param source_file: Image.open() 得到的 Image 对象
    :param new_width: 目标宽度（可选）
    :param new_height: 目标高度（可选）
    :param resample: NEAREST / BILINEAR / BICUBIC / LANCZOS
    :param ref_file: 参考图片路径（等比适配到参考尺寸）
    :return: PIL.Image (RGBA)
    """

    img = source_file.convert("RGBA")
    src_w, src_h = img.size

    # -------------------------
    # 计算目标尺寸（等比）
    # -------------------------
    if ref_file:
        ref_img = Image.open(ref_file)
        ref_w, ref_h = ref_img.size

        scale = min(ref_w / src_w, ref_h / src_h)
        target_w = int(src_w * scale)
        target_h = int(src_h * scale)

    else:
        if new_width > 0 and new_height > 0:
            scale = min(new_width / src_w, new_height / src_h)
            target_w = int(src_w * scale)
            target_h = int(src_h * scale)

        elif new_width > 0:
            scale = new_width / src_w
            target_w = new_width
            target_h = int(src_h * scale)

        elif new_height > 0:
            scale = new_height / src_h
            target_h = new_height
            target_w = int(src_w * scale)

        else:
            # 未指定任何尺寸，直接返回
            return img

    # -------------------------
    # 重采样方式
    # -------------------------
    resample_map = {
        "NEAREST": Image.Resampling.NEAREST,
        "BILINEAR": Image.Resampling.BILINEAR,
        "BICUBIC": Image.Resampling.BICUBIC,
        "LANCZOS": Image.Resampling.LANCZOS,
    }
    resample_method = resample_map.get(resample, Image.Resampling.LANCZOS)

    # -------------------------
    # Alpha 安全 resize（逐通道）
    # -------------------------
    bands = img.split()
    resized_bands = [
        band.resize((target_w, target_h), resample=resample_method) for band in bands
    ]

    return Image.merge("RGBA", resized_bands)


# 图片粘贴
def image_paste(paste_image, under_image, pos):
    """

    :param paste_image: 需要粘贴的图片
    :param under_image: 底图
    :param pos: 位置（x,y）坐标
    :return: 返回图片
    """
    # 获取需要贴入图片的透明通道
    r, g, b, alpha = paste_image.split()
    # 粘贴时将alpha值传递至mask属性
    under_image.paste(paste_image, pos, alpha)
    return under_image


def image_paste_center(paste_image, under_image, pos):
    """
    将一张图片粘贴到另一张图片上，并使粘贴图片的Y轴中心对齐到指定y坐标

    :param paste_image: 需要粘贴的图片 (RGBA)
    :param under_image: 底图 (RGBA)
    :param pos: 粘贴位置 (x, y)，其中y表示目标对齐的中心位置
    :return: 返回合成后的图片
    """
    x, y = pos
    # 计算上下居中偏移
    paste_w, paste_h = paste_image.size
    y_top = y - paste_h // 2  # 上边位置 = 中心点y - 半高

    # 获取透明通道
    r, g, b, alpha = paste_image.split()
    under_image.paste(paste_image, (x, y_top), alpha)
    return under_image


# 下载QQ头像
def download_avatar(user_id: str) -> bytes | None:
    url = f"http://q1.qlogo.cn/g?b=qq&nk={user_id}&s=640"
    data = download_url(url, timeout=10)
    if not data or hashlib.md5(data).hexdigest() == QQ_DEFAULT_AVATAR_MD5:
        url = f"http://q1.qlogo.cn/g?b=qq&nk={user_id}&s=100"
        data = download_url(url, timeout=10)
    return data


# 根据URL下载文件
def download_url(url: str, timeout=None) -> bytes | None:
    for i in range(3):
        try:
            resp = requests.get(url, timeout=timeout)
            if resp.status_code != 200:
                continue
            return resp.content
        except Exception as e:
            logger.warning(f"Error downloading {url}, retry {i}/3: {str(e)}")


# QQ 无头像时 qlogo 返回的默认图，md5 固定
QQ_DEFAULT_AVATAR_MD5 = "acef72340ac0e914090bd35799f5594e"


# 校验字节流能不能当图片打开
def _is_image(data) -> bool:
    try:
        Image.open(BytesIO(data)).size
        return True
    except Exception:
        return False


# 按QQ号获取头像字节流：先读本地缓存，未命中则下载并落盘
def get_avatar_bytes(user_id):
    if not user_id:
        return None
    cache_file = AVATAR_CACHE_DIR / f"{user_id}.png"
    if cache_file.exists():
        try:
            cached = cache_file.read_bytes()
            if _is_image(cached):
                return cached
            # 缓存已损坏（旧版本写入或写了一半），删掉重下，避免一直读到坏图
            logger.warning(f"头像缓存已损坏，重新下载：{cache_file}")
            cache_file.unlink()
        except Exception as err:
            logger.error(f"读取头像缓存失败：{str(err)}")

    try:
        data = download_avatar(str(user_id))
    except Exception as err:
        logger.error(f"下载头像失败：{str(err)}")
        return None
    if not data:
        logger.warning(f"未获取到QQ {user_id} 的头像数据")
        return None

    # 确认是图片再缓存，避免把错误页写进缓存
    if not _is_image(data):
        logger.error(f"头像数据不是有效图片，QQ：{user_id}")
        return None

    try:
        AVATAR_CACHE_DIR.mkdir(parents=True, exist_ok=True)
        cache_file.write_bytes(data)
        logger.info(f"头像已缓存至：{cache_file}")
    except Exception as err:
        logger.error(f"写入头像缓存失败：{str(err)}")
    return data


# 灰色圆形占位头像（BGRA，圆外透明）
def placeholder_avatar(size: int):
    avatar = np.zeros((size, size, 4), dtype=np.uint8)
    center = size // 2
    cv2.circle(avatar, (center, center), size // 2, (90, 90, 90, 255), -1)
    cv2.circle(
        avatar, (center, int(size * 0.38)), int(size * 0.16), (150, 150, 150, 255), -1
    )
    cv2.ellipse(
        avatar,
        (center, int(size * 0.96)),
        (int(size * 0.28), int(size * 0.30)),
        0,
        180,
        360,
        (150, 150, 150, 255),
        -1,
    )
    return avatar


# 载入指定尺寸的圆形头像，统一返回BGRA四通道
def load_avatar(user_id, size: int):
    data = get_avatar_bytes(user_id)
    avatar = (
        cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_UNCHANGED)
        if data
        else None
    )
    if avatar is not None and avatar.size:
        if avatar.ndim == 2:  # 灰度图
            avatar = cv2.cvtColor(avatar, cv2.COLOR_GRAY2BGRA)
        elif avatar.shape[2] == 3:
            avatar = cv2.cvtColor(avatar, cv2.COLOR_BGR2BGRA)
        elif avatar.shape[2] != 4:
            logger.warning(f"头像通道数异常（{avatar.shape[2]}），改用占位头像")
            avatar = None
    else:
        avatar = None

    if avatar is None:
        logger.warning(f"未获取到头像（QQ：{user_id}），使用占位头像")
        avatar = placeholder_avatar(size)
    return cv2.resize(avatar, (size, size), interpolation=cv2.INTER_AREA)


def draw_centered_text(draw, text, x_left, x_right, y, font, fill):
    """
    在指定区域水平居中绘制文本
    :param draw: ImageDraw对象
    :param text: 要绘制的文本
    :param x_left: 区域左边界x坐标
    :param x_right: 区域右边界x坐标
    :param y: 文本基线y坐标
    :param font: ImageFont对象
    :param fill: 文本颜色
    """
    text = str(text)
    text_width = font.getlength(text)  # 获取文本像素宽度
    area_width = x_right - x_left
    x_center = x_left + (area_width - text_width) / 2
    draw.text((x_center, y), text, font=font, fill=fill)


def draw_right_aligned_text(draw, text, x_left, x_right, y, font, fill):
    """
    在指定区域内从右往左绘制文本（右对齐）
    :param draw: ImageDraw 对象
    :param text: 要绘制的文本
    :param x_left: 区域左边界 x
    :param x_right: 区域右边界 x（文本右侧贴齐这里）
    :param y: 文本基线 y 坐标
    :param font: ImageFont 对象
    :param fill: 文本颜色
    """
    text = str(text)
    text_width = font.getlength(text)

    # 右对齐：文本右边贴到 x_right
    x_start = x_right - text_width

    # 可选：防止越界（超出左边界时截断）
    if x_start < x_left:
        x_start = x_left

    draw.text((x_start, y), text, font=font, fill=fill)


def get_save_icon(game, name, icon_type, url):
    if game == "bf6":
        name = name.replace("/", "").upper()
    else:
        name = name.replace("/", "")
    logger.info(f"查询的物品: {name},图标类型：{icon_type},游戏：{game}")
    icon = get_icon_from_cache(name, game, icon_type)
    if not icon:
        if game == "bf1":
            icon = get_icon_from_url(url)
            if icon_type == "weapon":
                icon = png_resize(icon, 500, 125)
            else:
                icon = png_resize(icon, 315, 79)
        elif game == "bf3":
            icon = get_icon_from_url(url)
            if icon_type == "weapon":
                icon = png_resize(icon, 500, 300)
            else:
                icon = png_resize(icon, 315, 200)
        elif game == "bf4":
            icon = get_icon_from_url(url)
            if icon_type == "weapon":
                icon = png_resize(icon, 500, 165)
            else:
                icon = png_resize(icon, 315, 80)
        elif game == "bfv":
            icon = get_icon_from_url(url)
            if icon_type == "weapon":
                icon = png_resize(icon, 500, 166)
            else:
                icon = png_resize(icon, 315, 315)
        # bf6 没有下载分支:缓存没命中时 icon 仍是 None,交给下面的占位图兜底
        if icon:
            icon_path = RES_DIR / "img" / game / icon_type / f"{name}.png"
            icon.save(icon_path)
            logger.info(f"文件已经缓存至：{icon_path}")

    # 上述步骤处理完再执行一次icon判空
    if icon:
        return icon
    else:
        # 针对空图片处理
        if icon_type == "weapon":
            icon = Image.open(RES_DIR / "img" / "wp.png").convert("RGBA")
        else:
            icon = Image.open(RES_DIR / "img" / "vc.png").convert("RGBA")
        return icon


# 获取icon
def get_icon_from_url(url):
    if url:
        try:
            # 添加10s超时判断，如果超时返回空
            res = BytesIO(requests.get(url, timeout=10).content)
            icon = Image.open(res).convert("RGBA")
            return icon
        except requests.exceptions.RequestException as e:
            logger.error(f"请求异常：{e}")


def get_icon_from_cache(icon_name, game, icon_type):
    path = RES_DIR / "img" / game / icon_type
    logger.info(f"查询物品：{icon_name}物品路径：{path}")
    try:
        icon_list = os.listdir(path)
        # logger.info(f"所有物品：{icon_list}")
        if icon_name in str(icon_list):
            logger.info(f"本地存在{icon_name}物品")
            img = Image.open(path / f"{icon_name}.png").convert("RGBA")
            return img
    except Exception as err:
        logger.error(f"获取图标失败：{str(err)}")
        return None


def cutout_region(img, mask, alpha=0):
    """
    在图像指定区域设置透明度
    :param img: 输入图像 (BGR 或 BGRA)
    :param mask: 区域定义
                 - 矩形: (x1, y1, x2, y2)
                 - 圆形: (cx, cy, r)
                 - 任意 mask: numpy.ndarray (0/255)
    :param alpha: 设置的透明度 (0=透明, 255=不透明, 0-255之间半透明)
    :return: 处理后的图像 (BGRA)
    """

    # 转为 BGRA (确保有 alpha 通道)
    if img.shape[2] == 3:
        b, g, r = cv2.split(img)
        a = np.ones(b.shape, dtype=np.uint8) * 255
        img = cv2.merge((b, g, r, a))

    h, w = img.shape[:2]
    region_mask = np.zeros((h, w), dtype=np.uint8)

    # 根据传入参数类型生成区域 mask
    if isinstance(mask, tuple):
        if len(mask) == 4:  # 矩形
            x1, y1, x2, y2 = mask
            region_mask[y1:y2, x1:x2] = 255
        elif len(mask) == 3:  # 圆形
            cx, cy, r = mask
            cv2.circle(region_mask, (cx, cy), r, 255, -1)
        else:
            raise ValueError("不支持的 mask tuple 格式")
    elif isinstance(mask, np.ndarray):  # 自定义 mask
        if mask.shape != (h, w):
            raise ValueError("mask 尺寸必须与图像一致")
        region_mask = mask
    else:
        raise TypeError("mask 必须是 tuple 或 numpy.ndarray")

    # 修改 alpha 通道
    img[region_mask > 0, 3] = alpha

    return img


def paste_image_cv2(fg, bg, pos):
    """
    将前景图粘贴到背景图上，支持透明度 (Alpha 混合)
    :param fg: 前景图 (必须是 BGRA)
    :param bg: 背景图 (BGR 或 BGRA)
    :param pos: 粘贴位置 (x, y) 左上角坐标
    :return: 合成后的图像 (BGRA)
    """

    x, y = pos
    fh, fw = fg.shape[:2]

    # 确保背景有 alpha 通道
    if bg.shape[2] == 3:
        b, g, r = cv2.split(bg)
        a = np.ones(b.shape, dtype=np.uint8) * 255
        bg = cv2.merge((b, g, r, a))

    # 截取 ROI 区域
    roi = bg[y : y + fh, x : x + fw]

    # 防止越界
    if roi.shape[0] != fh or roi.shape[1] != fw:
        raise ValueError("前景图超出背景范围")

    # 提取 alpha 通道
    alpha_fg = fg[:, :, 3] / 255.0
    alpha_bg = 1.0 - alpha_fg

    # 混合 RGB
    for c in range(3):
        roi[:, :, c] = alpha_fg * fg[:, :, c] + alpha_bg * roi[:, :, c]

    # 更新 alpha 通道（取最大值，保证透明度正确）
    roi[:, :, 3] = np.clip(fg[:, :, 3] + roi[:, :, 3], 0, 255)

    # 放回背景
    bg[y : y + fh, x : x + fw] = roi
    return bg


def notice_paste(cv2_bg):
    """贴公告栏（notice.png）到 BF1/3/4/5 卡片右上。

    公告图是 RGBA（透明底）素材：必须按 alpha 混合，否则透明区域会以白底
    整块盖在卡片上。没有 alpha 通道的旧素材仍走「三角切角 + 不透明覆盖」老流程。
    """
    pos = (1176, 318)
    front_img = cv2.imread(str(RES_DIR / "notice" / "notice.png"), cv2.IMREAD_UNCHANGED)
    if front_img is None:
        logger.warning("公告栏图片缺失（res/notice/notice.png），跳过公告栏绘制")
        return cv2_bg

    if front_img.shape[2] == 4:
        return _paste_rgba(front_img, cv2_bg, pos)

    # 旧素材（无 alpha）：三角切角后按不透明贴图
    if front_img.shape[:2] != (631, 1355):
        logger.warning(
            f"公告栏图片尺寸 {front_img.shape[:2]} 与切角蒙版 (631, 1355) 不一致，跳过公告栏绘制"
        )
        return cv2_bg
    mask = np.zeros((631, 1355), dtype=np.uint8)
    pts = np.array([[0, 0], [167, 0], [0, 631]], np.int32)
    cv2.fillPoly(mask, [pts], 255)
    front_img = cutout_region(front_img, mask, alpha=0)
    return paste_image_cv2(front_img, cv2_bg, pos)


def _paste_rgba(front_img, cv2_bg, pos):
    """RGBA 前景按 alpha 混合到背景（超出背景的部分自动裁剪）。"""
    x, y = pos
    h, w = front_img.shape[:2]

    # ROI 边界裁剪
    bg_h, bg_w = cv2_bg.shape[:2]
    w = min(w, bg_w - x)
    h = min(h, bg_h - y)

    roi = cv2_bg[y : y + h, x : x + w]

    alpha = front_img[:h, :w, 3] / 255.0
    alpha = alpha[:, :, None]

    cv2_bg[y : y + h, x : x + w] = (
        alpha * front_img[:h, :w, :3] + (1 - alpha) * roi
    ).astype(np.uint8)
    return cv2_bg


def notice_paste_bf6(cv2_bg):
    notice_path = RES_DIR / "notice" / "notice_bf6.png"
    front_img = cv2.imread(str(notice_path), cv2.IMREAD_UNCHANGED)
    if front_img is None:
        logger.warning(f"公告栏图片缺失（{notice_path}），跳过公告栏绘制")
        return cv2_bg

    if front_img.shape[2] == 4:
        # 有 Alpha 通道 → Alpha 混合
        return _paste_rgba(front_img, cv2_bg, (0, 0))

    # 无 Alpha → 直接覆盖（保留边界裁剪）
    bg_h, bg_w = cv2_bg.shape[:2]
    h = min(front_img.shape[0], bg_h)
    w = min(front_img.shape[1], bg_w)
    cv2_bg[0:h, 0:w] = front_img[:h, :w]
    return cv2_bg


def designation_paste(
    player, cv2_bg, user_id, start_x=10, start_y=10, target_height=40, spacing=6
):
    designation_list_path = get_designation_list_by_user_id(user_id, player)
    if not designation_list_path:
        return cv2_bg

    icons = []
    total_width = 0

    for path in designation_list_path:
        icon = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        if icon is None:
            continue

        h, w = icon.shape[:2]
        scale = target_height / h
        new_w = int(w * scale)

        icon = cv2.resize(icon, (new_w, target_height))
        icons.append(icon)
        total_width += new_w

    total_width += spacing * (len(icons) - 1)

    x = start_x
    y = start_y  # 关键变化点：直接使用 top-left

    bg_h, bg_w = cv2_bg.shape[:2]

    for icon in icons:
        ih, iw = icon.shape[:2]

        if x + iw > bg_w or y + ih > bg_h:
            break

        roi = cv2_bg[y : y + ih, x : x + iw]

        if icon.shape[2] == 4:
            alpha = icon[:, :, 3] / 255.0
            for c in range(3):
                roi[:, :, c] = roi[:, :, c] * (1 - alpha) + icon[:, :, c] * alpha
        else:
            roi[:] = icon

        cv2_bg[y : y + ih, x : x + iw] = roi
        x += iw + spacing

    return cv2_bg


def get_designation_list_by_user_id(user_id, player):
    # 标签是装饰性内容，查库失败（缺表、库损坏等）不该拖垮整张卡片，返回空即可
    try:
        usermanager = UserManager()
        designationmanager = DesignationManager()
        user_info = usermanager.get_user_by_qq(user_id)
        icon_no_list = []
        if (
            user_info is not None
            and user_info["ea_player_name"].upper() == player.upper()
        ):
            logger.info(f"玩家标签列表：{user_info}")
            icon_no_list = user_info["designation"]
        icon_list = designationmanager.get_designation_by_id_list(icon_no_list)
        icon_path_list = []
        for icon in icon_list:
            icon_file_name = icon["icon_path"]
            icon_path_list.append(
                str(RES_DIR / "img" / "icon" / "designation" / icon_file_name)
            )
        return icon_path_list
    except Exception as err:
        logger.warning(f"获取玩家标签失败（跳过标签绘制）：{str(err)}")
        return []


def paste_rank_icon(rank, img):
    rank_pic_path = get_rank_pic(rank)
    rank_icon_path = RES_DIR / "img" / "icon" / "rank" / rank_pic_path
    rank_icon = Image.open(rank_icon_path)
    rank_icon = png_resize(rank_icon, new_width=132, new_height=192)
    fm_img = image_paste(rank_icon, img, (946, 61))
    return fm_img


def get_rank_pic(rank: int) -> str:
    """
    根据 rank 返回对应的 rank_pic 路径
    """
    for (start, end), info in rank_pic.items():
        if start <= rank <= end:
            return info["path"]
    return "t_ui_rankswatch_1-9 1.png"