avatar.py

"""头像定位 + 运行时自学习（MAA 方案）。

完整流程（移植自 MAA BattlefieldMatcher + BattleHelper::update_deployment_）：
1. detect_slots: BattleOpersFlag 模板匹配 → 所有干员槽位
2. locate_oper: 遍历槽位 → 有缓存则 TemplateMatch → 无缓存则点击+OCR+存头像
3. OCR ROI 来自 MAA BattleOperName task: [5, 177, 191, 37]
"""

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np

if TYPE_CHECKING:
    from maa.context import Context
    from maa.controller import Controller

from aao.utils.logger import logger

# MAA 常量（1280×720）
_DETAIL_WAIT = 1  # 等详情页打开


def _avatar_dir() -> Path:
    from aao.utils.runtime_paths import project_root

    d = project_root() / "resource" / "base" / "image" / "avatar"
    d.mkdir(parents=True, exist_ok=True)
    return d


def _match_avatar_roi_offset() -> tuple[int, int, int, int]:
    """读取 reco.json 中 MatchAvatar.roi_offset（flag_rect → avatar_rect）。"""
    from aao.utils.runtime_paths import project_root

    path = project_root() / "resource" / "base" / "pipeline" / "reco.json"
    raw = json.loads(path.read_text(encoding="utf-8"))
    offset = raw["MatchAvatar"].get("roi_offset", [0, 0, 0, 0])
    return tuple(int(v) for v in offset)  # type: ignore[return-value]


def _apply_roi_offset(rect: tuple[int, int, int, int]) -> tuple[int, int, int, int]:
    ox, oy, ow, oh = _match_avatar_roi_offset()
    x, y, w, h = rect
    return x + ox, y + oy, w + ow, h + oh


def _get_char_id(oper_name: str) -> str:
    from aao.utils.runtime_paths import project_root

    mapping_path = project_root() / "data" / "operator_mapping.json"
    if not mapping_path.exists():
        return ""
    mapping = json.loads(mapping_path.read_text(encoding="utf-8"))
    return mapping.get(oper_name, "")


def _normalize_name(name: str) -> str:
    """干员名匹配前的最小清理。

    PP-OCRv6 rec 模型足够强，不再做 NFKC/字形替换/大小写折叠，避免把
    特殊干员名（如 GALLUS²）改写成与资源名不一致的形式。
    """
    return name.strip() if name else ""


def detect_slots(
    context: Context,
    image: np.ndarray,
) -> list[dict[str, Any]]:
    """用 pipeline 节点 DetectSlots 检测待部署区所有干员槽位。"""
    reco_detail = context.run_recognition("DetectSlots", image)

    if not reco_detail or not reco_detail.hit:
        return []

    slots = []
    for result in reco_detail.all_results:
        box = getattr(result, "box", None)
        if box is None:
            continue
        fx, fy, fw, fh = (int(v) for v in box)
        flag_rect = (fx, fy, fw, fh)
        avatar_rect = _apply_roi_offset(flag_rect)
        click_x = int(fx - 45 + 75 // 2)
        click_y = int(fy + 6 + 120 // 2)
        slots.append(
            {
                "flag_rect": flag_rect,
                "avatar_rect": avatar_rect,
                "click_pos": (click_x, click_y),
            }
        )

    logger.info("检测到 %d 个干员槽位", len(slots))
    return slots


def has_avatar(oper_name: str) -> bool:
    char_id = _get_char_id(oper_name)
    if not char_id:
        return False
    return (_avatar_dir() / f"{char_id}.png").exists()


def locate_oper(
    context: Context,
    ctrl: Controller,
    oper_name: str,
) -> tuple[float, float] | None:
    """定位指定干员在待部署区的位置。

    MAA 方案：
    1. 检测所有槽位
    2. 有缓存 → TemplateMatch 匹配
    3. 无缓存/匹配失败 → 逐个点击槽位 → OCR 干员名 → 截取头像存盘
    """
    image = ctrl.post_screencap().wait().get()
    if image is None:  # pyright: ignore[reportUnnecessaryComparison]
        logger.error("截图失败, 游戏可能已退出或最小化")
        return None
    slots = detect_slots(context, image)
    if not slots:
        logger.error("未检测到干员槽位")
        return None

    char_id = _get_char_id(oper_name)
    h, w = image.shape[:2]

    # Step 1: 有缓存 → 在每个槽位做 TemplateMatch
    if char_id and (_avatar_dir() / f"{char_id}.png").exists():
        for i, slot in enumerate(slots):
            ax, ay, aw, ah = slot["avatar_rect"]
            # MatchAvatar 的 roi_offset 写在 pipeline/reco.json 中；这里动态传 flag_rect。
            roi = list(slot["flag_rect"])

            reco = context.run_recognition(
                "MatchAvatar",
                image,
                pipeline_override={
                    "MatchAvatar": {
                        "template": f"avatar/{char_id}.png",
                        "roi": roi,
                    }
                },
            )

            if reco and reco.hit:
                cx = ax + aw // 2
                cy = ay + ah // 2
                logger.info("干员 %s 在槽位 %d", oper_name, i)
                return (cx / w, cy / h)

        logger.info("干员 %s 有缓存但未匹配，转入 OCR 学习", oper_name)

    # Step 2: 无缓存或匹配失败 → 点击每个未识别槽位 → OCR → 存头像
    for i, slot in enumerate(slots):
        click_x, click_y = slot["click_pos"]

        # 点击打开详情页
        logger.debug("点击槽位 %d (%d, %d)", i, click_x, click_y)
        ctrl.post_click(click_x, click_y).wait()
        time.sleep(_DETAIL_WAIT)

        # 截图详情页
        detail_img = ctrl.post_screencap().wait().get()
        if detail_img is None:  # pyright: ignore[reportUnnecessaryComparison]
            logger.warning("截图失败, 跳过该槽位")
            continue

        # OCR 干员名
        name = _ocr_oper_name(context, detail_img)
        logger.info("槽位 %d OCR: %s", i, name or "(空)")

        # 关闭详情页（再点一次）
        ctrl.post_click(click_x, click_y).wait()
        time.sleep(0.3)

        if name:
            matched = _normalize_name(name) == _normalize_name(oper_name)
            # 匹配则用原名存（保证 char_id 正确），否则用 OCR 名存（供后续复用）
            _save_avatar_from_image(image, slot, oper_name if matched else name)

            if matched:
                logger.info("找到目标干员 %s 在槽位 %d（OCR=%s）", oper_name, i, name)
                return (click_x / w, click_y / h)

    logger.error("未找到干员 %s", oper_name)
    return None


def _ocr_oper_name(context: Context, detail_img: np.ndarray) -> str | None:
    """OCR 读取详情页干员名。"""
    reco = context.run_recognition("OcrOperName", detail_img)

    if not reco or not reco.hit:
        return None

    # OCR 文字在 result.text 里（OCRResult 继承 BoxAndScoreResult + text）
    best = reco.best_result
    if best is not None:
        text = getattr(best, "text", "")
        if text:
            return text.strip()

    for result in reco.all_results:
        text = getattr(result, "text", "")
        if text:
            return text.strip()

    return None


def _save_avatar_from_image(
    image: np.ndarray,
    slot: dict[str, Any],
    oper_name: str,
) -> bool:
    """从截图截取槽位头像并存盘。"""
    char_id = _get_char_id(oper_name)
    if not char_id:
        return False

    ax, ay, aw, ah = slot["avatar_rect"]
    h, w = image.shape[:2]
    x1, y1 = max(0, ax), max(0, ay)
    x2, y2 = min(w, ax + aw), min(h, ay + ah)

    avatar = image[y1:y2, x1:x2]
    if avatar.size == 0:
        return False

    from PIL import Image

    out_path = _avatar_dir() / f"{char_id}.png"
    avatar_rgb = avatar[..., ::-1].copy()
    Image.fromarray(avatar_rgb).save(out_path)
    logger.info("头像已缓存: %s → %s", oper_name, out_path)
    return True


def list_cached() -> list[str]:
    return sorted(p.name for p in _avatar_dir().glob("*.png"))