HeXi/hexi/plugins/nonebot_plugin_helldivers_tools/handlers/war.py

"""绝地潜兵小助手 —— 触发器(Trigger)层。"""

from __future__ import annotations

import os
import re
from pathlib import Path
from typing import Optional

import httpx
from nonebot import logger, on_command
from nonebot.adapters.onebot.v11 import MessageEvent, MessageSegment
from nonebot.internal.params import ArgPlainText
from nonebot.matcher import Matcher
from PIL import Image

from ..services.equipment import get_equipment_by_combination, get_random_equipment
from ..services.hd2_api import get_briefing_data
from ..utils import gen_ms_img, pic2b64
from ..utils.war_renderer import render_war_briefing

_BASE_DIR = Path(__file__).resolve().parents[1]
_IMG_DIR = _BASE_DIR / "res" / "img"


async def _send_war_card(matcher: Matcher, ev: MessageEvent, top_per_race: Optional[int]) -> None:
    await matcher.send("正在获取前线战况，请民主的等待！")
    try:
        data = await get_briefing_data()
        png = await render_war_briefing(data, top_per_race)
    except Exception as e:
        logger.exception("获取战况失败")
        await matcher.finish(
            MessageSegment.reply(ev.message_id)
            + f"战况获取失败：{type(e).__name__}：{e}"
        )
        return
    mes = (MessageSegment.reply(ev.message_id), gen_ms_img(png))
    await matcher.send(mes)


war_situation = on_command("简报", aliases={"简报"})
war_detail = on_command("详报", aliases={"详报"})


@war_situation.handle()
async def get_war_info(ev: MessageEvent):
    await _send_war_card(war_situation, ev, top_per_race=4)


@war_detail.handle()
async def get_war_detail(ev: MessageEvent):
    await _send_war_card(war_detail, ev, top_per_race=None)


async def download_url(url: str) -> bytes | None:
    async with httpx.AsyncClient() as client:
        for i in range(3):
            try:
                resp = await client.get(url)
                if resp.status_code != 200:
                    continue
                return resp.content
            except Exception as e:
                logger.error(f"Error downloading {url}, retry {i}/3: {str(e)}")
        return None


random_helldivers = on_command("随机战备", aliases={"随机战备"}, block=True)

board_path = str(_IMG_DIR / "board2.png")
PROMPT = MessageSegment.image(pic2b64(Image.open(board_path)))

error_count = 0


@random_helldivers.got("pick_type", prompt=PROMPT)
async def got_random_helldivers(event: MessageEvent, pick_type: str = ArgPlainText()):
    global error_count
    if error_count >= 2:
        await random_helldivers.finish("已回归平民生活——")
        logger.info("持续输错退出随机战备")
        error_count = 0
        return
    logger.info(f"用户选择的战备类型: {pick_type}")
    if not is_number(pick_type):
        error_count += 1
        await random_helldivers.reject(f"您输入的 {pick_type} 非数字，请重新输入1到11，或者输入0退出")
    elif int(pick_type) not in range(12):
        error_count += 1
        await random_helldivers.reject(f"您输入的 {pick_type} 不在范围内，请重新输入1到11，或者输入0退出")
    elif int(pick_type) == 0:
        await random_helldivers.finish("已回归平民生活——")
        logger.info("用户选择退出随机战备")
        error_count = 0
        return

    mix_msg = (MessageSegment.reply(event.message_id),)

    type_combinations = {
        2: {"red": 2, "blue": 1, "green": 1},
        3: {"green": 2, "red": 1, "blue": 1},
        4: {"blue": 2, "green": 1, "red": 1},
        5: {"red": 3, "blue": 1},
        6: {"green": 3, "blue": 1},
        7: {"red": 2, "blue": 2},
        8: {"blue": 2, "green": 2},
        9: {"green": 2, "red": 2},
        10: {"red": 4},
        11: {"green": 4},
    }

    if int(pick_type) == 1:
        logger.info("用户选择纯随机战备")
        result = await get_random_equipment(4)
    else:
        combination = type_combinations.get(int(pick_type))
        logger.info(f"用户选择的战备组合: {combination}")
        result = await get_equipment_by_combination(combination)

    final_msg = MessageSegment.text("您的随机结果是：\n")
    image_tuple = gen_ms_img(result)
    mix_msg += (final_msg, image_tuple)

    await random_helldivers.finish(mix_msg)


def is_number(s):
    return bool(re.match(r"^[0-9]+$", s))