"""
rust_fmt_fixer.py  v6.0
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
修复 v5.0 的两个核心 Bug:

Bug1: total_size 越界 —— 把相邻独立字符串也纳入范围
  原因: _compute_total_size 的扫描逻辑贪婪地读取后续可打印字节
  修复: 严格限制 total_size = Decoder 解析停止位置，
        不再向后扩展，不包含末尾 \0

Bug2: make_array 覆盖范围超出 fmt 模板
  原因: IDA make_array(ea, n) 中 n 是元素数，
        total_size 错误导致 n 过大
  修复: 精确传入 total_size，并在定义前验证边界

设计原则:
  - fmt 模板字节数组: [base_ea, base_ea + total_size)
  - 末尾 \0 (tag=0 的空字符串) 包含在内
  - 不触碰 base_ea + total_size 之后的任何字节
"""

import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple

import idc
import idaapi
import idautils
import ida_bytes
import ida_ida
import ida_nalt


# ═══════════════════════════════════════════════════════════════
#  §1  配置
# ═══════════════════════════════════════════════════════════════

class Cfg:
    verbose      : bool  = True
    dry_run      : bool  = False
    min_str_len  : int   = 1
    max_str_len  : int   = 8192
    min_printable: float = 0.55
    lea_scan_back: int   = 64
    max_pieces   : int   = 64
    add_comments : bool  = True
    auto_rename  : bool  = True
    max_cmt_len  : int   = 512

    AUTO_PREFIXES: Tuple[str, ...] = (
        "stru_", "byte_", "unk_", "word_",
        "dword_", "qword_", "off_", "asc_", "sub_", "loc_",
    )
    FMT_NEW_SIGS: Tuple[str, ...] = (
        "_ZN4core3fmt9Arguments3new",
        "_ZN4core3fmt9Arguments16new_v1",
        "fmt::Arguments::new",
        "core::fmt::Arguments::new",
        "Arguments3new",
    )


# ═══════════════════════════════════════════════════════════════
#  §2  日志
# ═══════════════════════════════════════════════════════════════

class Log:
    @staticmethod
    def _pr(tag: str, msg: str) -> None:
        print(f"[RustFmt][{tag}] {msg}")

    @classmethod
    def i(cls, m: str)  -> None: cls._pr("INFO ", m)
    @classmethod
    def w(cls, m: str)  -> None: cls._pr("WARN ", m)
    @classmethod
    def e(cls, m: str)  -> None: cls._pr("ERROR", m)
    @classmethod
    def ok(cls, m: str) -> None: cls._pr("OK   ", m)
    @classmethod
    def d(cls, m: str)  -> None:
        if Cfg.verbose:
            cls._pr("DBG  ", m)

    @staticmethod
    def sec(title: str) -> None:
        bar = "─" * 58
        print(f"\n┌{bar}┐")
        print(f"│  {title:<56}│")
        print(f"└{bar}┘")


# ═══════════════════════════════════════════════════════════════
#  §3  IDA 工具层
# ═══════════════════════════════════════════════════════════════

class IDA:

    @staticmethod
    def is_64bit() -> bool:
        try:
            if hasattr(ida_ida, "inf_is_64bit"):
                return ida_ida.inf_is_64bit()
            return idaapi.get_inf_structure().is_64bit()
        except Exception:
            return True

    PTR: int = 8

    @classmethod
    def read_ptr(cls, ea: int) -> Optional[int]:
        try:
            if not ida_bytes.is_loaded(ea):
                return None
            if cls.PTR == 8:
                return ida_bytes.get_qword(ea) & 0xFFFF_FFFF_FFFF_FFFF
            return ida_bytes.get_dword(ea) & 0xFFFF_FFFF
        except Exception:
            return None

    @staticmethod
    def read_bytes(ea: int, n: int) -> Optional[bytes]:
        try:
            if n <= 0 or n > 0x10_0000:
                return None
            if not ida_bytes.is_loaded(ea):
                return None
            return ida_bytes.get_bytes(ea, n)
        except Exception:
            return None

    @staticmethod
    def read_u8(ea: int) -> Optional[int]:
        try:
            if not ida_bytes.is_loaded(ea):
                return None
            return ida_bytes.get_byte(ea)
        except Exception:
            return None

    @staticmethod
    def is_loaded(ea: int) -> bool:
        if ea in (0, idc.BADADDR) or ea is None:
            return False
        try:
            return bool(ida_bytes.is_loaded(ea))
        except Exception:
            return False

    @staticmethod
    def seg_name(ea: int) -> str:
        return idc.get_segm_name(ea) or ""

    @classmethod
    def in_rodata(cls, ea: int) -> bool:
        n = cls.seg_name(ea).lower()
        return any(k in n for k in
                   (".rdata", ".rodata", ".data.rel.ro", "__const"))

    @staticmethod
    def get_name(ea: int) -> str:
        return idc.get_name(ea, idc.GN_VISIBLE) or ""

    @staticmethod
    def set_name(ea: int, name: str) -> bool:
        try:
            return bool(idc.set_name(
                ea, name,
                idc.SN_NOWARN | idc.SN_NOCHECK | idc.SN_FORCE,
            ))
        except Exception:
            return False

    @staticmethod
    def all_names() -> List[Tuple[int, str]]:
        out: List[Tuple[int, str]] = []
        try:
            for item in idautils.Names():
                if not (isinstance(item, (list, tuple)) and len(item) == 2):
                    continue
                a, b = item
                if isinstance(a, int) and isinstance(b, str):
                    out.append((a, b))
                elif isinstance(a, str) and isinstance(b, int):
                    out.append((b, a))
        except Exception as ex:
            Log.w(f"all_names: {ex}")
        return out

    @staticmethod
    def set_cmt(ea: int, cmt: str) -> None:
        try:
            idc.set_cmt(ea, cmt[:Cfg.max_cmt_len], 0)
        except Exception:
            pass

    @staticmethod
    def get_cmt(ea: int) -> str:
        try:
            return idc.get_cmt(ea, 0) or ""
        except Exception:
            return ""

    @staticmethod
    def append_cmt(ea: int, text: str) -> None:
        old = IDA.get_cmt(ea)
        if text in old:
            return
        IDA.set_cmt(ea, f"{old} | {text}" if old else text)

    # ── 数据定义（精确范围控制）────────────────────────

    @staticmethod
    def force_undef(ea: int, size: int) -> None:
        """
        仅清除 [ea, ea+size) 范围内的定义。
        使用 DELIT_SIMPLE 而非 DELIT_EXPAND，
        避免意外扩展到相邻数据项。
        """
        try:
            # DELIT_SIMPLE: 只删除 ea 处的项，不扩展
            ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)
        except Exception:
            pass

    @staticmethod
    def force_undef_exact(ea: int, size: int) -> None:
        """
        逐字节清除，确保精确边界控制。
        用于 fmt 模板块的清除，不影响 ea+size 之后的数据。
        """
        # 先尝试整体清除
        try:
            ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)
        except Exception:
            pass

        # 再逐字节确认清除（处理跨越大项目的情况）
        offset = 0
        while offset < size:
            cur_ea = ea + offset
            try:
                item_sz = max(1, ida_bytes.get_item_size(cur_ea))
                ida_bytes.del_items(cur_ea, ida_bytes.DELIT_SIMPLE, 1)
                offset += 1
            except Exception:
                offset += 1

    @staticmethod
    def make_byte_array(ea: int, size: int) -> bool:
        """
        将 [ea, ea+size) 定义为字节数组。

        !! 关键: size 就是字节数，不是元素数 !!
        IDA make_array(ea, count) 中 count = 元素个数
        对于 byte 类型，元素大小=1，所以 count = size
        """
        if size <= 0:
            return False
        try:
            # 清除旧定义（精确范围）
            ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)

            # 定义第一个字节
            if not idc.create_byte(ea):
                return False

            # 扩展为数组（count = 字节数）
            if size > 1:
                idc.make_array(ea, size)

            # 验证：检查实际定义的大小
            actual = ida_bytes.get_item_size(ea)
            if actual != size:
                Log.d(f"    make_byte_array: 请求{size}字节，实际{actual}字节")
                # 不报错，继续（IDA 有时会合并相邻项）

            return True
        except Exception as ex:
            Log.d(f"    make_byte_array 异常: {ex}")
            return False

    @staticmethod
    def make_qword(ea: int) -> bool:
        try:
            ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, 8)
            return bool(idc.create_qword(ea))
        except Exception:
            return False

    @staticmethod
    def make_string(ea: int, length: int) -> bool:
        try:
            ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length + 1)
            return bool(ida_bytes.create_strlit(ea, length, ida_nalt.STRTYPE_C))
        except Exception:
            return False

    @staticmethod
    def get_flags(ea: int) -> int:
        try:
            return ida_bytes.get_full_flags(ea)
        except Exception:
            return 0

    @staticmethod
    def is_qword(ea: int) -> bool:
        return ida_bytes.is_qword(IDA.get_flags(ea))

    @staticmethod
    def is_strlit(ea: int) -> bool:
        return ida_bytes.is_strlit(IDA.get_flags(ea))

    @staticmethod
    def get_item_size(ea: int) -> int:
        try:
            return ida_bytes.get_item_size(ea)
        except Exception:
            return 1

    @staticmethod
    def has_xref(ea: int) -> bool:
        """检查 ea 处是否有数据交叉引用（即被外部引用）"""
        try:
            for _ in idautils.XrefsTo(ea):
                return True
        except Exception:
            pass
        return False

    @staticmethod
    def mnem(ea: int) -> str:
        return idc.print_insn_mnem(ea).lower()

    @staticmethod
    def operand(ea: int, n: int) -> str:
        return idc.print_operand(ea, n).lower()

    @staticmethod
    def op_type(ea: int, n: int) -> int:
        return idc.get_operand_type(ea, n)

    @staticmethod
    def op_val(ea: int, n: int) -> int:
        return idc.get_operand_value(ea, n)

    @staticmethod
    def prev_head(ea: int) -> int:
        return idc.prev_head(ea)

    @staticmethod
    def xrefs_to(ea: int) -> list:
        try:
            return list(idautils.XrefsTo(ea))
        except Exception:
            return []

    @staticmethod
    def xrefs_from(ea: int) -> list:
        try:
            return list(idautils.XrefsFrom(ea))
        except Exception:
            return []

    @staticmethod
    def refresh() -> None:
        try:
            idaapi.refresh_idaview_anyway()
        except Exception:
            pass


IDA.PTR = 8 if IDA.is_64bit() else 4


# ═══════════════════════════════════════════════════════════════
#  §4  字符串工具
# ═══════════════════════════════════════════════════════════════

def printable_ratio(data: bytes) -> float:
    if not data:
        return 1.0
    good = sum(1 for b in data if 0x20 <= b <= 0x7E or b in (9, 10, 13))
    return good / len(data)


def decode_utf8(ea: int, length: int) -> Optional[str]:
    if not (Cfg.min_str_len <= length <= Cfg.max_str_len):
        return None
    if not IDA.is_loaded(ea):
        return None
    raw = IDA.read_bytes(ea, length)
    if raw is None:
        return None
    try:
        text = raw.decode("utf-8", errors="replace")
        good = sum(1 for c in text if c.isprintable() or c in "\n\r\t")
        if len(text) > 0 and (good / len(text)) < Cfg.min_printable:
            return None
        return text
    except Exception:
        return None


def make_id(s: str, maxlen: int = 28) -> str:
    r = re.sub(r"[^a-zA-Z0-9]", "_", s[:maxlen])
    r = re.sub(r"_+", "_", r).strip("_").lower()
    return r or "fmt"


def is_auto_name(ea: int) -> bool:
    name = IDA.get_name(ea)
    return bool(name) and any(name.startswith(p) for p in Cfg.AUTO_PREFIXES)


# ═══════════════════════════════════════════════════════════════
#  §5  数据结构
# ═══════════════════════════════════════════════════════════════

@dataclass
class StrPiece:
    ea            : int
    length        : int
    content       : str
    is_placeholder: bool = False
    tag_ea        : int  = 0


@dataclass
class DecodeResult:
    fmt_name  : str
    pieces    : List[StrPiece]
    raw_ea    : int
    # !! 关键字段: 解码器严格记录的解析字节数
    # 不包含任何后续字节，精确等于 TLV 序列的长度
    total_size: int = 0
    extra     : Dict = field(default_factory=dict)

    @property
    def strings(self) -> List[str]:
        return [p.content for p in self.pieces if not p.is_placeholder]

    @property
    def human_readable(self) -> str:
        return "".join(
            "{}" if p.is_placeholder else p.content
            for p in self.pieces
        )

    @property
    def preview(self) -> str:
        hr = self.human_readable
        return hr[:80] if hr else "<empty>"


# ═══════════════════════════════════════════════════════════════
#  §6  解码器（精确 total_size）
# ═══════════════════════════════════════════════════════════════

class Decoder:

    # ── FORMAT_A ────────────────────────────────────────

    def try_format_a(self, ea: int) -> Optional[DecodeResult]:
        P = IDA.PTR
        pieces_ptr = IDA.read_ptr(ea)
        pieces_len = IDA.read_ptr(ea + P)

        if pieces_ptr is None or pieces_len is None:
            return None
        if not IDA.is_loaded(pieces_ptr):
            return None
        if not (0 < pieces_len <= Cfg.max_pieces):
            return None
        if not IDA.in_rodata(pieces_ptr):
            return None

        pieces: List[StrPiece] = []
        for i in range(pieces_len):
            slot    = pieces_ptr + i * P * 2
            str_ptr = IDA.read_ptr(slot)
            str_len = IDA.read_ptr(slot + P)
            if str_ptr is None or str_len is None:
                break
            if str_ptr == 0 and str_len == 0:
                pieces.append(StrPiece(0, 0, ""))
                continue
            if not IDA.is_loaded(str_ptr):
                break
            if str_len > Cfg.max_str_len:
                break
            content = decode_utf8(str_ptr, str_len)
            if content is None:
                break
            pieces.append(StrPiece(str_ptr, str_len, content))

        if not pieces:
            return None
        return DecodeResult("FORMAT_A", pieces, ea)

    # ── FORMAT_B ─────────────────────────────────────────

    def try_format_b(self, ea: int) -> Optional[DecodeResult]:
        """
        严格 TLV 解析，精确记录 total_size。

        终止条件（任意一个满足即停止）：
          1. 遇到 tag=0 (空字符串，通常是序列结束标记)
          2. 连续解析失败次数 > 阈值
          3. 超过最大扫描字节数

        total_size = 解析停止时的 offset（从 ea 起的字节偏移）
        这个值精确界定了 fmt 模板的边界，不包含后续任何数据。
        """
        pieces          : List[StrPiece] = []
        offset          : int            = 0
        max_scan        : int            = 512   # 单个模板不超过 512 字节
        consecutive_fail: int            = 0
        MAX_FAIL        : int            = 3

        while offset < max_scan:
            tag_ea   = ea + offset
            tag_byte = IDA.read_u8(tag_ea)

            if tag_byte is None:
                break

            # ── tag=0: 序列结束标记 ──────────────────────
            # Rust fmt 模板以 tag=0 (空字符串) 结尾
            if tag_byte == 0:
                pieces.append(StrPiece(
                    ea=ea + offset + 1, length=0, content="",
                    tag_ea=tag_ea,
                ))
                offset += 1  # 包含这个结束字节
                break        # ← 严格在此终止，不继续向后

            # ── 普通字符串片段: tag < 0x80 ───────────────
            if (tag_byte & 0x80) == 0:
                length  = tag_byte
                str_ea  = ea + offset + 1

                content = decode_utf8(str_ea, length)
                if content is None:
                    consecutive_fail += 1
                    if consecutive_fail >= MAX_FAIL:
                        break
                    offset += 1
                    continue

                pieces.append(StrPiece(
                    ea=str_ea, length=length, content=content,
                    tag_ea=tag_ea,
                ))
                Log.d(f"    B str +{offset:#05x} len={length} '{content[:40]}'")
                offset += 1 + length
                consecutive_fail = 0

            # ── 占位符: tag & 0xC0 == 0xC0 ───────────────
            elif (tag_byte & 0xC0) == 0xC0:
                pieces.append(StrPiece(
                    ea=0, length=0, content="{}",
                    is_placeholder=True, tag_ea=tag_ea,
                ))
                Log.d(f"    B ph  +{offset:#05x} tag={tag_byte:#04x}")
                offset += 1
                consecutive_fail = 0

            # ── 其他高字节 ───────────────────────────────
            else:
                Log.d(f"    B ctl +{offset:#05x} {tag_byte:#04x}")
                consecutive_fail += 1
                if consecutive_fail >= MAX_FAIL:
                    break
                offset += 1

        real = [p for p in pieces if not p.is_placeholder and p.length > 0]
        if not real:
            return None

        result = DecodeResult("FORMAT_B", pieces, ea)
        # !! total_size 严格等于解析消耗的字节数
        # 这就是 fmt 模板的精确大小，后续数组定义不得超出此范围
        result.total_size = offset
        Log.d(f"    B total_size={offset} (0x{ea:X}..0x{ea+offset:X})")
        return result

    # ── FORMAT_C ─────────────────────────────────────────

    def try_format_c(self, ea: int) -> Optional[DecodeResult]:
        raw = IDA.read_bytes(ea, 256)
        if raw is None:
            return None

        pieces    : List[StrPiece] = []
        cur       : bytearray      = bytearray()
        cur_start : int            = 0

        for i, b in enumerate(raw):
            if 0x20 <= b <= 0x7E or b in (9, 10, 13):
                if not cur:
                    cur_start = i
                cur.append(b)
            else:
                if len(cur) >= 3:
                    try:
                        text = cur.decode("utf-8", errors="strict")
                        pieces.append(StrPiece(
                            ea=ea + cur_start, length=len(cur), content=text,
                        ))
                    except Exception:
                        pass
                cur = bytearray()

        if len(cur) >= 3:
            try:
                pieces.append(StrPiece(
                    ea=ea + cur_start, length=len(cur),
                    content=cur.decode("utf-8", errors="strict"),
                ))
            except Exception:
                pass

        if not pieces:
            return None
        return DecodeResult("FORMAT_C", pieces, ea)

    # ── 主入口 ───────────────────────────────────────────

    def decode(self, ea: int) -> Optional[DecodeResult]:
        Log.d(f"  decode 0x{ea:X}")
        for tag, method in (
            ("A", self.try_format_a),
            ("B", self.try_format_b),
            ("C", self.try_format_c),
        ):
            r = method(ea)
            if r:
                Log.d(f"  ✓ {tag}: {r.preview}")
                return r
            Log.d(f"  ✗ {tag}")
        return None


# ═══════════════════════════════════════════════════════════════
#  §7  调用点追踪
# ═══════════════════════════════════════════════════════════════

_ARG0_REGS = ("rcx", "ecx", "rdi", "edi")


@dataclass
class CallSite:
    call_ea    : int
    func_ea    : int
    template_ea: int
    result     : Optional[DecodeResult]

    @property
    def strings(self) -> List[str]:
        return self.result.strings if self.result else []

    @property
    def preview(self) -> str:
        return self.result.preview if self.result else "<no decode>"


class CallTracer:

    def __init__(self):
        self._decoder = Decoder()

    def trace(self, call_ea: int) -> Optional[CallSite]:
        func_ea = idc.get_func_attr(call_ea, idc.FUNCATTR_START)
        tmpl_ea = self._find_arg0(call_ea)
        if tmpl_ea is None:
            return None
        result = self._decoder.decode(tmpl_ea)
        return CallSite(
            call_ea=call_ea, func_ea=func_ea,
            template_ea=tmpl_ea, result=result,
        )

    def _find_arg0(self, call_ea: int) -> Optional[int]:
        cur     = call_ea
        reg_map : Dict[str, int] = {}

        for _ in range(Cfg.lea_scan_back):
            cur = IDA.prev_head(cur)
            if cur == idc.BADADDR:
                break

            m   = IDA.mnem(cur)
            op0 = IDA.operand(cur, 0)

            if m == "lea":
                addr = self._mem_op_ea(cur, 1)
                if addr is not None and IDA.is_loaded(addr):
                    reg_map[op0] = addr
                    if op0 in _ARG0_REGS:
                        return addr

            elif m == "mov" and op0 in _ARG0_REGS:
                op1 = IDA.operand(cur, 1).strip()
                if op1 in reg_map:
                    return reg_map[op1]
                addr = self._imm_op_ea(cur, 1)
                if addr is not None and IDA.is_loaded(addr):
                    return addr

            for xref in IDA.xrefs_from(cur):
                xtype = getattr(xref, "type", -1)
                if xtype in (1, 2, 3, 4, 16, 17, 18, 19, 20):
                    ref_ea = xref.to
                    if IDA.is_loaded(ref_ea) and IDA.in_rodata(ref_ea):
                        reg_map[op0] = ref_ea
                        if op0 in _ARG0_REGS:
                            return ref_ea

        return None

    @staticmethod
    def _mem_op_ea(insn_ea: int, op_idx: int) -> Optional[int]:
        ot = IDA.op_type(insn_ea, op_idx)
        ov = IDA.op_val(insn_ea, op_idx)
        if ot in (idc.o_mem, idc.o_displ):
            return ov
        if ot == idc.o_imm and IDA.is_loaded(ov):
            return ov
        return None

    @staticmethod
    def _imm_op_ea(insn_ea: int, op_idx: int) -> Optional[int]:
        ot = IDA.op_type(insn_ea, op_idx)
        ov = IDA.op_val(insn_ea, op_idx)
        if ot in (idc.o_imm, idc.o_mem) and IDA.is_loaded(ov):
            return ov
        return None


# ═══════════════════════════════════════════════════════════════
#  §8  扫描器
# ═══════════════════════════════════════════════════════════════

class Scanner:

    def __init__(self):
        self._tracer = CallTracer()

    def find_fmt_functions(self) -> List[int]:
        found: Set[int] = set()
        for ea, name in IDA.all_names():
            for sig in Cfg.FMT_NEW_SIGS:
                if sig in name:
                    found.add(ea)
                    Log.i(f"  ✓ {name} @ 0x{ea:X}")
                    break
        return sorted(found)

    def scan(self) -> List[CallSite]:
        fmt_fns = self.find_fmt_functions()
        if not fmt_fns:
            Log.w("未找到 fmt::Arguments::new")
            return []

        sites      : List[CallSite] = []
        seen_calls : Set[int]       = set()

        for fn_ea in fmt_fns:
            refs = IDA.xrefs_to(fn_ea)
            Log.i(f"  函数 0x{fn_ea:X}: {len(refs)} 个引用")

            for xref in refs:
                call_ea = xref.frm
                if call_ea in seen_calls:
                    continue
                seen_calls.add(call_ea)

                if IDA.mnem(call_ea) not in ("call", "jmp", "bl", "blx"):
                    continue

                site = self._tracer.trace(call_ea)
                if site:
                    sites.append(site)

        return sites


# ═══════════════════════════════════════════════════════════════
#  §9  修复器
# ═══════════════════════════════════════════════════════════════

class Fixer:

    def __init__(self):
        self.n_sites   = 0
        self.n_fixed   = 0
        self.n_renamed = 0
        self.n_retyped = 0
        self.n_errors  = 0
        # 记录已处理的模板地址，防止重复修复破坏相邻数据
        self._done: Set[int] = set()

    # ── 公共入口 ─────────────────────────────────────────

    def fix_site(self, site: CallSite) -> bool:
        self.n_sites += 1
        try:
            return self._fix(site)
        except Exception as ex:
            Log.e(f"fix_site 0x{site.call_ea:X}: {ex}")
            self.n_errors += 1
            return False

    def fix_result(self, ea: int, result: DecodeResult) -> bool:
        try:
            return self._apply(ea, result, call_ea=None)
        except Exception as ex:
            Log.e(f"fix_result 0x{ea:X}: {ex}")
            return False

    # ── 内部实现 ─────────────────────────────────────────

    def _fix(self, site: CallSite) -> bool:
        if site.result is None:
            if Cfg.add_comments and not Cfg.dry_run:
                IDA.append_cmt(site.call_ea, "fmt::Arguments::new")
            return False
        ok = self._apply(site.template_ea, site.result, site.call_ea)
        if ok:
            self.n_fixed += 1
        return ok

    def _apply(self,
               data_ea : int,
               result  : DecodeResult,
               call_ea : Optional[int]) -> bool:

        # 同一模板地址只处理一次（可能被多个 call 引用）
        if data_ea in self._done:
            if call_ea and Cfg.add_comments and not Cfg.dry_run:
                IDA.append_cmt(call_ea, f'fmt!("{result.human_readable}")')
            return True
        self._done.add(data_ea)

        if Cfg.dry_run:
            Log.ok(f"[DRY] 0x{data_ea:X} [{result.fmt_name}]"
                   f" size={result.total_size}"
                   f' "{result.preview}"')
            return True

        # ① call 注释
        if call_ea and Cfg.add_comments:
            IDA.append_cmt(call_ea, f'fmt!("{result.human_readable}")')

        # ② 数据类型修复
        if result.fmt_name == "FORMAT_A":
            self._fix_format_a(data_ea, result)
        elif result.fmt_name == "FORMAT_B":
            self._fix_format_b(data_ea, result)

        # ③ 起始地址注释
        if Cfg.add_comments:
            IDA.set_cmt(
                data_ea,
                f'const core::fmt::Arguments\n'
                f'fmt!("{result.human_readable}")',
            )

        # ④ 重命名
        if result.strings and Cfg.auto_rename:
            safe = make_id(result.strings[0])
            if is_auto_name(data_ea):
                if IDA.set_name(data_ea, f"rust_fmt_{safe}"):
                    self.n_renamed += 1

        Log.ok(f'0x{data_ea:X} [{result.fmt_name}]'
               f' size={result.total_size}'
               f' "{result.preview}"')
        return True

    # ── FORMAT_A ────────────────────────────────────────

    def _fix_format_a(self, base_ea: int, result: DecodeResult) -> None:
        P = IDA.PTR
        pieces_ptr = IDA.read_ptr(base_ea)
        if pieces_ptr is None:
            return

        for i, piece in enumerate(result.pieces):
            if piece.length == 0:
                continue
            slot_ea = pieces_ptr + i * P * 2
            ida_bytes.del_items(slot_ea, ida_bytes.DELIT_SIMPLE, P * 2)
            IDA.make_qword(slot_ea)
            IDA.make_qword(slot_ea + P)
            IDA.set_cmt(slot_ea, f'&str[{i}]: "{piece.content[:40]}"')

            if IDA.is_loaded(piece.ea) and piece.length > 0:
                if not IDA.is_strlit(piece.ea):
                    if IDA.make_string(piece.ea, piece.length):
                        self.n_retyped += 1
                IDA.append_cmt(piece.ea, f'utf8: "{piece.content[:40]}"')
                if is_auto_name(piece.ea):
                    safe = make_id(piece.content)
                    if IDA.set_name(piece.ea, f"rust_str_{safe}"):
                        self.n_renamed += 1

    # ── FORMAT_B（核心修复）──────────────────────────────

    def _fix_format_b(self, base_ea: int, result: DecodeResult) -> None:
        """
        将 [base_ea, base_ea + total_size) 定义为整体字节数组。

        !! 安全保证 !!
          1. 仅操作 total_size 字节，不超出边界
          2. 操作前检查边界处是否有外部 xref 引用
             如果有，说明 total_size 计算有误，缩短到该处
          3. 使用 DELIT_SIMPLE 而非 DELIT_EXPAND
        """
        total = result.total_size
        if total <= 0:
            Log.w(f"    FORMAT_B: total_size=0 @ 0x{base_ea:X}，跳过")
            return

        # 安全边界检查：确认 [base_ea, base_ea+total) 内没有外部引用
        # 外部引用意味着该地址是独立数据的起始点，不该被覆盖
        safe_total = self._safe_size(base_ea, total)
        if safe_total != total:
            Log.w(f"    边界收缩: {total} → {safe_total} "
                  f"(0x{base_ea+safe_total:X} 有外部引用)")
        total = safe_total

        if total <= 0:
            Log.w(f"    FORMAT_B: 安全大小=0，跳过 0x{base_ea:X}")
            return

        Log.d(f"    FORMAT_B: make_byte_array 0x{base_ea:X} [{total}]")

        # 清除旧定义（精确范围，不扩展）
        ida_bytes.del_items(base_ea, ida_bytes.DELIT_SIMPLE, total)

        # 定义为整体字节数组
        ok = self._create_byte_array(base_ea, total)
        if ok:
            self.n_retyped += 1
        else:
            Log.w(f"    make_byte_array 失败 @ 0x{base_ea:X}[{total}]")

    def _safe_size(self, base_ea: int, requested: int) -> int:
        """
        从 base_ea+1 开始，找到第一个有外部 xref 的地址。
        返回从 base_ea 到该地址的字节数（不超过 requested）。

        base_ea 本身有 xref 是正常的（这是 fmt 模板的入口引用）。
        """
        for offset in range(1, requested):
            check_ea = base_ea + offset
            if not IDA.is_loaded(check_ea):
                return offset

            # 检查是否有外部数据引用
            for xref in IDA.xrefs_to(check_ea):
                xtype = getattr(xref, "type", -1)
                # 数据引用类型: dr_O(1), dr_W(2), dr_R(3), dr_T(4)
                if xtype in (1, 2, 3, 4):
                    Log.d(f"      xref @ +{offset} (0x{check_ea:X})"
                          f" from 0x{xref.frm:X} type={xtype}")
                    return offset

        return requested

    @staticmethod
    def _create_byte_array(ea: int, size: int) -> bool:
        """
        创建字节数组的最终实现。

        为什么不用 make_array:
          idc.make_array(ea, count) 有时会将 count 解释为
          "创建 count 个与 ea 处相同类型的元素"，
          行为依赖 IDA 版本，不稳定。

        更可靠的方式:
          使用 ida_bytes.create_data 直接创建指定大小的数据项。
        """
        try:
            # 方式1: create_data (最可靠)
            # FF_BYTE = 0x00000000 (字节类型标志)
            ok = ida_bytes.create_data(
                ea,
                ida_bytes.FF_BYTE,
                size,
                idc.BADADDR,
            )
            if ok:
                return True
        except Exception as ex:
            Log.d(f"    create_data: {ex}")

        try:
            # 方式2: create_byte + make_array
            if idc.create_byte(ea):
                if size > 1:
                    idc.make_array(ea, size)
                return True
        except Exception as ex:
            Log.d(f"    create_byte+make_array: {ex}")

        return False

    # ── 报告 ─────────────────────────────────────────────

    def report(self) -> None:
        Log.sec("执行报告")
        Log.i(f"  处理调用点数 : {self.n_sites}")
        Log.i(f"  成功修复数   : {self.n_fixed}")
        Log.i(f"  数据重定义数 : {self.n_retyped}")
        Log.i(f"  符号重命名数 : {self.n_renamed}")
        Log.i(f"  错误数       : {self.n_errors}")
        if self.n_errors == 0:
            Log.ok("全部完成")
        Log.i("提示: F5 刷新 | Edit→Undo 撤销")


# ═══════════════════════════════════════════════════════════════
#  §10  主引擎
# ═══════════════════════════════════════════════════════════════

class Engine:

    def __init__(self):
        self._scanner = Scanner()
        self._fixer   = Fixer()

    def run(self) -> None:
        Log.sec("Rust fmt! 修复工具 v6.0  精确边界控制")
        Log.i(f"IDA  : {idaapi.get_kernel_version()}")
        Log.i(f"架构 : {'64-bit' if IDA.is_64bit() else '32-bit'}")
        Log.i(f"DRY  : {Cfg.dry_run}")

        Log.sec("阶段1  扫描")
        sites = self._scanner.scan()
        Log.i(f"调用点: {len(sites)}")
        if not sites:
            Log.w("无调用点 | 手动: fix_call(0x...) / fix_data(0x...)")
            return

        Log.sec("阶段2  修复")
        total = len(sites)
        for idx, site in enumerate(sites, 1):
            self._fixer.fix_site(site)
            if idx % 50 == 0 or idx == total:
                Log.i(f"  {idx}/{total}")

        IDA.refresh()
        self._fixer.report()


# ═══════════════════════════════════════════════════════════════
#  §11  便捷 API
# ═══════════════════════════════════════════════════════════════

_tracer  = CallTracer()
_decoder = Decoder()
_fixer   = Fixer()


def fix_call(call_ea: int) -> None:
    """通过 call 指令地址修复"""
    Log.sec(f"fix_call 0x{call_ea:X}")
    site = _tracer.trace(call_ea)
    if site:
        _fixer.fix_site(site)
    else:
        Log.w("追踪失败，尝试 fix_data()")


def fix_data(data_ea: int) -> None:
    """直接修复数据地址"""
    Log.sec(f"fix_data 0x{data_ea:X}")
    result = _decoder.decode(data_ea)
    if result:
        _fixer.fix_result(data_ea, result)
    else:
        Log.w("解码失败，尝试 dump()")


def dump(ea: int, size: int = 96) -> None:
    """打印字节并解码（不修改 IDA）"""
    raw = IDA.read_bytes(ea, size)
    if raw is None:
        print(f"无法读取 0x{ea:X}")
        return

    print(f"\n━━ 0x{ea:X} ({size} bytes) ━━")
    for i in range(0, len(raw), 16):
        chunk = raw[i:i+16]
        hex_s = " ".join(f"{b:02X}" for b in chunk)
        asc_s = "".join(chr(b) if 0x20 <= b <= 0x7E else "·" for b in chunk)
        print(f"  +{i:04X}  {hex_s:<48}  {asc_s}")

    print("\n━━ 解码 ━━")
    old, Cfg.verbose = Cfg.verbose, True
    result = _decoder.decode(ea)
    Cfg.verbose = old

    if result:
        print(f"  格式      : {result.fmt_name}")
        print(f"  总大小    : {result.total_size} bytes  "
              f"(0x{ea:X} → 0x{ea+result.total_size:X})")
        print(f'  格式字符串: "{result.human_readable}"')
        print(f"  片段:")
        for i, p in enumerate(result.pieces):
            kind = "PH " if p.is_placeholder else "STR"
            print(f"    [{i:02d}] {kind} tag=0x{p.tag_ea:X}"
                  f" data=0x{p.ea:X} len={p.length}"
                  f" '{p.content[:40]}'")

        # 显示边界之后的数据（验证不越界）
        end_ea  = ea + result.total_size
        next_raw = IDA.read_bytes(end_ea, 16)
        if next_raw:
            hex_s = " ".join(f"{b:02X}" for b in next_raw)
            asc_s = "".join(chr(b) if 0x20 <= b <= 0x7E else "·"
                            for b in next_raw)
            print(f"\n  边界后 0x{end_ea:X}: {hex_s}  {asc_s}")
            print(f"  (此处数据应不被修改)")
    else:
        print("  解码失败")


def verify(ea: int) -> None:
    """验证修复结果：检查 ea 处的数据类型和相邻数据完整性"""
    result = _decoder.decode(ea)
    if result is None:
        print(f"0x{ea:X}: 解码失败")
        return

    total = result.total_size
    end   = ea + total
    size  = IDA.get_item_size(ea)

    print(f"\n━━ verify 0x{ea:X} ━━")
    print(f"  预期大小 : {total} bytes  → 0x{end:X}")
    print(f"  实际大小 : {size} bytes")
    print(f"  大小匹配 : {'✓' if size == total else '✗'}")

    # 检查边界后的数据是否完整
    print(f"\n  边界后 0x{end:X} 的数据:")
    for offset in range(0, 32, 8):
        chunk_ea = end + offset
        raw = IDA.read_bytes(chunk_ea, 8)
        if raw is None:
            break
        item_sz = IDA.get_item_size(chunk_ea)
        name    = IDA.get_name(chunk_ea)
        print(f"    0x{chunk_ea:X} [{item_sz:2d}B] {raw.hex()}"
              f"  {name or '(unnamed)'}")


def set_verbose(v: bool) -> None: Cfg.verbose = v
def set_dry_run(v: bool) -> None: Cfg.dry_run = v; Log.i(f"dry_run={v}")


if __name__ == "__main__":
    Engine().run()