"""
AutoPOE2 - Ghidra Static RE Extractor (TASK-RE-PIPELINE-2026 / Phase C2-C3)

Trích xuất tri thức Reverse Engineering TĨNH từ PathOfExile.exe bằng PyGhidra,
rồi xuất ra JSON để nạp vào Offset Registry (data/poe2_offsets.db / offsets.toml).

Thiết kế bám sát 2 invariant của kế hoạch Phương án C:
  INV-SIG-UNIQUE    : một AOB pattern chỉ được chấp nhận khi khớp ĐÚNG 1 lần
                      trong section .text. >=2 lần -> tự loại (không đưa vào registry).
                      Đây chính là lý do 4 pattern trong docs/development/06 mục 3
                      bị [MISS]/vô nghĩa trên client thật.
  INV-DATA-PROV     : mọi record đều kèm provenance (build_id, sha256, thời điểm,
                      công cụ, đường dẫn) theo Rule 14 - không dùng dữ liệu tự chế.

Chiến lược thu hẹp không gian tìm kiếm (không cần disassemble thủ công 48.77MB):
  1. Ghidra tự dựng function table từ .pdata (RUNTIME_FUNCTION) -> ~154.515 hàm.
  2. Quét toàn bộ reference có nguồn trong .text và đích trong .data/.rdata.
  3. Xếp hạng ứng viên theo FAN-IN (số instruction cùng trỏ vào một ô nhớ):
     global pointer kiểu GameController/InGameState luôn có fan-in cao.
  4. Với mỗi ứng viên, sinh AOB pattern từ chính instruction RIP-relative
     (wildcard 4 byte disp32 cuối lệnh) rồi kiểm chứng INV-SIG-UNIQUE.

Quy ước RIP offset khớp 100% với C++ AobScanner::ResolveRipTarget(reader, match, 3):
    pattern "48 8B 05 ? ? ? ?"  -> độ dài 7, disp32 tại offset 3 = length - 4.

Cách dùng:
  # Smoke test trên binary nhỏ (chứng minh pipeline chạy được, không cần game):
  python tools/re_pipeline/re_extract.py --exe bin/Release/AutoPOE2_MemProbe.exe

  # Chạy thật trên client (LÂU: auto-analysis .text 48.77MB có thể 1-6 giờ):
  python tools/re_pipeline/re_extract.py --max-candidates 40

Xem: docs/development/25_static_reverse_engineering_pipeline_and_offset_registry.md
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple

REPO_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(Path(__file__).resolve().parent))

import fingerprint as fp  # noqa: E402

DEFAULT_PROJECT_DIR = REPO_ROOT / "data" / "re_ghidra"
DEFAULT_OUT_DIR = REPO_ROOT / "data" / "re_out"

# Section chứa mã lệnh (nguồn của reference) và section chứa dữ liệu (đích)
CODE_BLOCK_HINTS = (".text",)
DATA_BLOCK_HINTS = (".data", ".rdata", "_RDATA")


def _java_bytes_to_bytes(java_array) -> bytes:
    """byte[] của Java có dấu -> bytes không dấu của Python."""
    return bytes(bytearray(b & 0xFF for b in java_array))


def count_pattern_occurrences(
    haystack: bytes, pattern: List[Optional[int]]
) -> int:
    """Đếm số lần khớp pattern có wildcard (None = wildcard) trong haystack.

    Tối ưu: tìm neo bằng byte cố định ĐẦU TIÊN rồi mới kiểm tra phần còn lại,
    tránh quét O(n*m) trên .text 48.77MB.
    """
    if not pattern:
        return 0

    anchor_index = next((i for i, b in enumerate(pattern) if b is not None), None)
    if anchor_index is None:
        return 0  # toàn wildcard -> vô nghĩa

    anchor_byte = bytes([pattern[anchor_index]])
    tail = pattern[anchor_index + 1 :]
    need_before = anchor_index

    count = 0
    start = 0
    total = len(haystack)
    while True:
        found = haystack.find(anchor_byte, start)
        if found < 0:
            break
        match_start = found - need_before
        if match_start >= 0 and match_start + len(pattern) <= total:
            ok = True
            for offset, expected in enumerate(tail):
                if expected is None:
                    continue
                if haystack[found + 1 + offset] != expected:
                    ok = False
                    break
            if ok:
                count += 1
        start = found + 1
    return count


def to_pattern_string(pattern: List[Optional[int]]) -> str:
    """Xuất pattern theo cú pháp AobPattern::Parse của C++ ('?' = wildcard)."""
    return " ".join("?" if b is None else f"{b:02X}" for b in pattern)


def extract_with_ghidra(
    exe_path: Path,
    project_dir: Path,
    analyze: bool,
    max_candidates: int,
    pattern_window: int,
) -> dict:
    """Mở binary trong Ghidra và trích xuất tri thức RE. Trả về dict kết quả."""
    import pyghidra  # noqa: PLC0415 - import muộn để --help không cần khởi động JVM

    install_dir = os.environ.get("GHIDRA_INSTALL_DIR")
    print(
        f"[ghidra] {datetime.now():%H:%M:%S} starting JVM "
        f"(GHIDRA_INSTALL_DIR={install_dir or '(auto)'})",
        flush=True,
    )
    pyghidra.start(install_dir=Path(install_dir) if install_dir else None)
    print(f"[ghidra] {datetime.now():%H:%M:%S} JVM ready", flush=True)

    from ghidra.program.model.address import AddressSet  # noqa: PLC0415

    project_dir.mkdir(parents=True, exist_ok=True)
    print(
        f"[ghidra] {datetime.now():%H:%M:%S} open_program "
        f"analyze={analyze} project={project_dir} "
        f"(auto-analysis .text ~48.77MB may take 1-6h)",
        flush=True,
    )

    with pyghidra.open_program(
        exe_path,
        project_location=project_dir,
        project_name="AutoPOE2_RE",
        analyze=analyze,
        program_name=exe_path.name,
    ) as flat_api:
        print(
            f"[ghidra] {datetime.now():%H:%M:%S} program open; "
            "scanning functions/references",
            flush=True,
        )
        program = flat_api.getCurrentProgram()
        memory = program.getMemory()
        listing = program.getListing()
        function_manager = program.getFunctionManager()
        reference_manager = program.getReferenceManager()
        address_factory = program.getAddressFactory()

        image_base = str(program.getImageBase())

        # ---- 1. Bảng memory block (đối chiếu section table đã fingerprint) ----
        blocks = []
        text_block = None
        data_blocks = []
        for block in memory.getBlocks():
            blocks.append(
                {
                    "name": block.getName(),
                    "start": str(block.getStart()),
                    "end": str(block.getEnd()),
                    "size": int(block.getSize()),
                    "execute": bool(block.isExecute()),
                    "write": bool(block.isWrite()),
                    "read": bool(block.isRead()),
                    "initialized": bool(block.isInitialized()),
                }
            )
            if not block.isInitialized():
                continue
            if block.isExecute() and text_block is None:
                text_block = block
            elif any(hint in block.getName() for hint in DATA_BLOCK_HINTS) or (
                block.isWrite() and not block.isExecute()
            ):
                data_blocks.append(block)

        if text_block is None:
            raise RuntimeError("Không tìm thấy block mã lệnh (.text) trong binary")

        # ---- 2. Function table (Ghidra dựng từ .pdata RUNTIME_FUNCTION) ----
        function_count = 0
        function_iterator = function_manager.getFunctions(True)
        while function_iterator.hasNext():
            function_iterator.next()
            function_count += 1
        print(
            f"[ghidra] {datetime.now():%H:%M:%S} function_count={function_count:,}",
            flush=True,
        )

        # ---- 3. Reference có nguồn trong .text, đích trong .data/.rdata ----
        # Fan-in = số instruction cùng trỏ vào một ô nhớ. Global pointer kiểu
        # GameController / InGameState luôn có fan-in cao bất thường.
        text_start = text_block.getStart()
        text_end = text_block.getEnd()
        text_range = AddressSet(text_start, text_end)

        data_ranges = [
            (block.getStart(), block.getEnd(), block.getName()) for block in data_blocks
        ]

        fan_in: Dict[str, dict] = {}
        scanned_references = 0
        # getReferenceSourceIterator() trả về ĐỊA CHỈ NGUỒN (Address), không phải
        # đối tượng Reference -> phải lấy Reference qua getReferencesFrom() cho
        # từng nguồn (bài học smoke test 10/09/2026: GenericAddress không có
        # getToAddress()).
        source_iterator = reference_manager.getReferenceSourceIterator(text_range, True)
        last_heartbeat = time.time()
        while source_iterator.hasNext():
            source_address = source_iterator.next()
            for reference in reference_manager.getReferencesFrom(source_address):
                scanned_references += 1
                now = time.time()
                if now - last_heartbeat >= 30.0:
                    print(
                        f"[ghidra] {datetime.now():%H:%M:%S} refs={scanned_references:,} "
                        f"slots={len(fan_in):,}",
                        flush=True,
                    )
                    last_heartbeat = now

                to_address = reference.getToAddress()
                owner = None
                for start, end, block_name in data_ranges:
                    if start <= to_address <= end:
                        owner = block_name
                        break
                if owner is None:
                    continue

                key = str(to_address)
                entry = fan_in.get(key)
                if entry is None:
                    entry = {"address": key, "block": owner, "fan_in": 0, "sources": []}
                    fan_in[key] = entry
                entry["fan_in"] += 1
                if len(entry["sources"]) < 8:
                    entry["sources"].append(str(source_address))

        # ---- 4. Sinh AOB pattern + kiểm chứng INV-SIG-UNIQUE ----
        # flat_api.getBytes(addr, length) là API chuẩn trả về byte[] Java
        # (MemoryMapDB.getBytes không có overload (Address, int) - smoke test 10/09/2026)
        text_bytes = _java_bytes_to_bytes(
            flat_api.getBytes(text_start, int(text_block.getSize()))
        )
        text_base_value = text_start.getOffset()

        print(
            f"[ghidra] {datetime.now():%H:%M:%S} refs done={scanned_references:,} "
            f"slots={len(fan_in):,}; generating AOB for top {max_candidates}",
            flush=True,
        )
        ranked = sorted(fan_in.values(), key=lambda e: e["fan_in"], reverse=True)
        candidates: List[dict] = []
        unique_signatures: List[dict] = []

        for entry in ranked[:max_candidates]:
            record = {
                "address": entry["address"],
                "block": entry["block"],
                "fan_in": entry["fan_in"],
                "sources": entry["sources"],
                "signature": None,
                "match_count": None,
                "unique": False,
            }

            for source_str in entry["sources"]:
                source_address = address_factory.getAddress(source_str)
                instruction = listing.getInstructionAt(source_address)
                if instruction is None:
                    continue

                length = instruction.getLength()
                # RIP-relative [rip+disp32]: disp32 luôn là 4 byte CUỐI của lệnh
                if length < 5 or length > pattern_window:
                    continue

                raw = _java_bytes_to_bytes(flat_api.getBytes(source_address, length))
                pattern = list(raw[: length - 4]) + [None, None, None, None]
                occurrences = count_pattern_occurrences(text_bytes, pattern)

                # INV-SIG-UNIQUE: nếu chưa unique, TỰ ĐỘNG MỞ RỘNG pattern bằng
                # các lệnh kế tiếp (giữ wildcard cho mọi disp32 RIP-relative) cho
                # tới khi khớp đúng 1 lần hoặc hết cửa sổ pattern_window.
                # Bằng chứng thực nghiệm (smoke test 10/09/2026): pattern ngắn
                # "FF 15 ?? ?? ??" khớp 677 lần, "48 8B 05 ?? ?? ??" khớp 57 lần
                # -> bắt buộc phải có ngữ cảnh mở rộng mới đạt unique.
                cursor = source_address
                while occurrences != 1:
                    next_instruction = listing.getInstructionAfter(cursor)
                    if next_instruction is None:
                        break
                    next_address = next_instruction.getAddress()
                    next_length = next_instruction.getLength()
                    if len(pattern) + next_length > pattern_window:
                        break

                    next_raw = list(
                        _java_bytes_to_bytes(flat_api.getBytes(next_address, next_length))
                    )
                    # Lệnh kế tiếp cũng RIP-relative -> wildcard 4 byte cuối
                    has_rip_ref = any(
                        reference.getReferenceType().isData()
                        for reference in reference_manager.getReferencesFrom(next_address)
                    )
                    if has_rip_ref and next_length >= 5:
                        for pad in range(4):
                            next_raw[next_length - 4 + pad] = None

                    pattern.extend(next_raw)
                    cursor = next_address
                    occurrences = count_pattern_occurrences(text_bytes, pattern)

                record["signature"] = to_pattern_string(pattern)
                record["match_count"] = occurrences
                record["unique"] = occurrences == 1
                record["rip_offset"] = length - 4
                record["instruction_address"] = source_str
                record["mnemonic"] = instruction.getMnemonicString()
                record["pattern_bytes"] = len(pattern)
                print(
                    f"[ghidra] {datetime.now():%H:%M:%S} candidate "
                    f"{len(candidates)+1}/{max_candidates} fan_in={entry['fan_in']} "
                    f"unique={record['unique']} matches={record['match_count']}",
                    flush=True,
                )
                break

            candidates.append(record)
            if record["unique"]:
                unique_signatures.append(record)

        return {
            "image_base": image_base,
            "memory_blocks": blocks,
            "function_count": function_count,
            "scanned_references": scanned_references,
            "text_block": {
                "name": text_block.getName(),
                "start": str(text_block.getStart()),
                "size": int(text_block.getSize()),
            },
            "distinct_global_slots": len(fan_in),
            "candidates": candidates,
            "unique_signatures": unique_signatures,
        }


def main() -> int:
    parser = argparse.ArgumentParser(
        description="AutoPOE2 Ghidra Static RE Extractor (Phase C2-C3)"
    )
    parser.add_argument(
        "--exe",
        default=None,
        help="Binary cần phân tích (mặc định: tự định vị PathOfExile.exe)",
    )
    parser.add_argument(
        "--project-dir",
        default=str(DEFAULT_PROJECT_DIR),
        help="Thư mục Ghidra project (mặc định: data/re_ghidra)",
    )
    parser.add_argument(
        "--out",
        default=None,
        help="Tệp JSON kết quả (mặc định: data/re_out/<BUILD_ID>/extraction.json)",
    )
    parser.add_argument(
        "--no-analyze",
        action="store_true",
        help="Bỏ qua auto-analysis (chỉ dùng khi project đã analyze trước đó)",
    )
    parser.add_argument(
        "--max-candidates",
        type=int,
        default=25,
        help="Số ứng viên global slot fan-in cao nhất sẽ sinh signature",
    )
    parser.add_argument(
        "--pattern-window",
        type=int,
        default=24,
        help="Độ dài tối đa (byte) của instruction được dùng làm AOB pattern",
    )
    args = parser.parse_args()

    exe_path = Path(args.exe) if args.exe else fp.locate_poe2_executable()
    if exe_path is None or not exe_path.is_file():
        print("[ERROR] Không xác định được binary cần phân tích.", file=sys.stderr)
        return 2

    print(f"[1/3] Fingerprint binary: {exe_path}")
    build = fp.load_pe_headers_with_hash(exe_path)
    if not build.valid:
        print(f"[ERROR] Parse PE thất bại: {build.error}", file=sys.stderr)
        return 1
    build_id = build.short_id
    print(
        f"      BuildId={build_id} TimeDateStamp=0x{build.time_date_stamp:08X} "
        f"Sections={build.number_of_sections} SHA256={build.file_sha256[:16]}..."
    )

    out_path = (
        Path(args.out)
        if args.out
        else DEFAULT_OUT_DIR / build_id / "extraction.json"
    )

    print(f"[2/3] Chạy Ghidra extraction (analyze={not args.no_analyze})...")
    print(
        "      Lưu ý: auto-analysis .text 48.77MB của PathOfExile.exe có thể mất 1-6 giờ."
    )
    started_at = time.time()
    try:
        result = extract_with_ghidra(
            exe_path,
            Path(args.project_dir),
            analyze=not args.no_analyze,
            max_candidates=args.max_candidates,
            pattern_window=args.pattern_window,
        )
    except Exception as exc:  # noqa: BLE001 - báo lỗi rõ ràng cho pipeline RE
        print(f"[ERROR] Ghidra extraction thất bại: {exc}", file=sys.stderr)
        return 1
    elapsed = time.time() - started_at

    # INV-DATA-PROV (Rule 14): mọi record phải truy vết được nguồn gốc thật
    payload = {
        "schema": "autopoe2.re_extraction/v1",
        "generated_at": datetime.now().strftime("%d/%m/%Y %H:%M:%S"),
        "elapsed_seconds": round(elapsed, 2),
        "tool": "PyGhidra (Ghidra 12.1.3)",
        "provenance": {
            "binary_path": str(exe_path),
            "build_id": build_id,
            "time_date_stamp": f"0x{build.time_date_stamp:08X}",
            "size_of_image": f"0x{build.size_of_image:X}",
            "file_size": build.file_size,
            "file_sha256": build.file_sha256,
            "section_layout_hash": f"0x{build.section_layout_hash():016X}",
        },
        "invariants": {
            "INV-SIG-UNIQUE": "signature chỉ được chấp nhận khi match_count == 1",
            "INV-BUILD-01": "mọi tri thức khóa theo build_id ở trên",
        },
        "result": result,
    }

    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(
        json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )

    print(f"[3/3] Đã ghi kết quả -> {out_path}  ({elapsed:.1f}s)")
    print(f"      Functions          : {result['function_count']:,}")
    print(f"      Scanned references : {result['scanned_references']:,}")
    print(f"      Global slots       : {result['distinct_global_slots']:,}")
    print(f"      Candidates         : {len(result['candidates'])}")
    print(f"      UNIQUE signatures  : {len(result['unique_signatures'])} (INV-SIG-UNIQUE)")
    for record in result["unique_signatures"][:10]:
        print(
            f"        fan_in={record['fan_in']:<5} {record['block']:<8} "
            f"{record['address']}  {record['signature']}"
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())