#!/usr/bin/env python3
"""Scan broad block write candidates that could cover opcode 0x24 mode1 source."""
from __future__ import annotations

import argparse
import html
import json
import struct
import sys
from pathlib import Path

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

from probe_exe_scene_tables import read_sections


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
MODE1_SOURCE = 0x0059E348
BLOCK_SCAN_START = 0x0059E300
BLOCK_SCAN_END = MODE1_SOURCE
WINDOW_BYTES = 48
MAX_BLOCK_SIZE = 0x2000
REG_NAMES = ["eax", "ecx", "edx", "ebx", "esp", "ebp", "esi", "edi"]
FAILED_OPCODE24_MODE1_BLOCK_WRITE_GATE_IDS = [
    "mode1-source-block-write",
    "covering-size-transfer-candidate",
    "runtime-watchpoint-trace",
]
OPCODE24_MODE1_BLOCK_WRITE_MISSING_EVIDENCE = [
    "broad block write candidate covering opcode 0x24 mode1 source 0x0059e348",
    "nearby covering size and transfer instruction proving a static block producer",
    "runtime watchpoint trace proving the 0x0059e348 producer",
]
OPCODE24_MODE1_BLOCK_WRITE_EVIDENCE_REFS = [
    {"path": "Hwanse2.exe", "description": "dword-reference block-write scan over 0x0059e300..0x0059e348"},
    {"path": "out/save_selector_opcode24_mode1_source_writes.json", "description": "direct/static source-write scan for mode1 source"},
    {"path": "out/runtime_trace_feasibility.json", "description": "runtime watchpoint trace availability blocker"},
]


def hex32(value: int) -> str:
    return f"0x{value:08x}"


def hex8(value: int) -> str:
    return f"0x{value:02x}"


def text_section(exe: bytes) -> tuple[dict, bytes]:
    section = next(section for section in read_sections(exe) if section["name"] == ".text")
    return section, exe[section["raw"] : section["raw"] + section["raw_size"]]


def find_dword_refs(data: bytes, target: int) -> list[int]:
    needle = struct.pack("<I", target)
    refs = []
    pos = data.find(needle)
    while pos >= 0:
        refs.append(pos)
        pos = data.find(needle, pos + 1)
    return refs


def classify_ref(data: bytes, pos: int) -> dict:
    if pos >= 6 and data[pos - 6] == 0xC7 and data[pos - 5] == 0x85:
        return {
            "accessKind": "address",
            "width": 4,
            "instruction": "mov dword ptr [ebp+disp32], imm32",
            "instructionKind": "mov-stack-imm32",
            "instructionPrefixBytes": 6,
        }
    if pos >= 3 and data[pos - 3] == 0xC7 and data[pos - 2] in {0x45, 0x44}:
        return {
            "accessKind": "address",
            "width": 4,
            "instruction": "mov dword ptr [ebp+disp8], imm32",
            "instructionKind": "mov-stack-imm32",
            "instructionPrefixBytes": 3,
        }
    if pos >= 2 and data[pos - 2 : pos] in {b"\x66\xa1", b"\x66\xa3"}:
        write = data[pos - 1] == 0xA3
        return {
            "accessKind": "write" if write else "read",
            "width": 2,
            "instruction": f"mov {'ds:[addr], ax' if write else 'ax, ds:[addr]'}",
            "instructionKind": "direct-moffs16",
            "instructionPrefixBytes": 2,
        }
    if pos >= 1:
        opcode = data[pos - 1]
        if opcode in {0xA0, 0xA1, 0xA2, 0xA3}:
            if opcode == 0xA0:
                return {
                    "accessKind": "read",
                    "width": 1,
                    "instruction": "mov al, ds:[addr]",
                    "instructionKind": "direct-moffs8",
                    "instructionPrefixBytes": 1,
                }
            if opcode == 0xA1:
                return {
                    "accessKind": "read",
                    "width": 4,
                    "instruction": "mov eax, ds:[addr]",
                    "instructionKind": "direct-moffs32",
                    "instructionPrefixBytes": 1,
                }
            if opcode == 0xA2:
                return {
                    "accessKind": "write",
                    "width": 1,
                    "instruction": "mov ds:[addr], al",
                    "instructionKind": "direct-moffs8",
                    "instructionPrefixBytes": 1,
                }
            return {
                "accessKind": "write",
                "width": 4,
                "instruction": "mov ds:[addr], eax",
                "instructionKind": "direct-moffs32",
                "instructionPrefixBytes": 1,
            }
        if 0xB8 <= opcode <= 0xBF:
            reg = REG_NAMES[opcode - 0xB8]
            return {
                "accessKind": "address",
                "width": 4,
                "instruction": f"mov {reg}, imm32",
                "instructionKind": "mov-reg-imm32",
                "instructionPrefixBytes": 1,
            }
        if opcode == 0x68:
            return {
                "accessKind": "address",
                "width": 4,
                "instruction": "push imm32",
                "instructionKind": "push-imm32",
                "instructionPrefixBytes": 1,
            }
    if pos >= 3 and data[pos - 3] == 0x66 and data[pos - 2] in {0x8B, 0x89}:
        modrm = data[pos - 1]
        if modrm & 0xC7 == 0x05:
            write = data[pos - 2] == 0x89
            return {
                "accessKind": "write" if write else "read",
                "width": 2,
                "instruction": f"opcode 0x66 0x{data[pos - 2]:02x} modrm 0x{modrm:02x} direct disp32",
                "instructionKind": "direct-modrm16",
                "instructionPrefixBytes": 3,
            }
    if pos >= 2 and data[pos - 2] in {0x8A, 0x8B, 0x88, 0x89, 0xC6, 0xC7, 0xFE, 0xFF, 0x80, 0x81, 0x83}:
        opcode = data[pos - 2]
        modrm = data[pos - 1]
        if modrm & 0xC7 == 0x05:
            write = opcode in {0x88, 0x89, 0xC6, 0xC7, 0xFE, 0xFF, 0x80, 0x81, 0x83}
            width = 1 if opcode in {0x8A, 0x88, 0xC6, 0xFE, 0x80, 0x83} else 4
            return {
                "accessKind": "write" if write else "read",
                "width": width,
                "instruction": f"opcode 0x{opcode:02x} modrm 0x{modrm:02x} direct disp32",
                "instructionKind": "direct-modrm",
                "instructionPrefixBytes": 2,
            }
    return {
        "accessKind": "unknown",
        "width": None,
        "instruction": "unclassified dword reference",
        "instructionKind": "unknown",
        "instructionPrefixBytes": 0,
    }


def small_immediates(data: bytes, start: int, end: int, text_va: int) -> list[dict]:
    rows = []
    index = start
    while index < min(end, len(data)):
        opcode = data[index]
        if opcode == 0x6A and index + 1 < len(data):
            value = struct.unpack_from("b", data, index + 1)[0]
            if 0 <= value <= MAX_BLOCK_SIZE:
                rows.append({
                    "vaHex": hex32(text_va + index),
                    "value": value,
                    "valueHex": hex8(value),
                    "instruction": "push imm8",
                })
            index += 2
            continue
        if opcode in {0x68, 0xB8, 0xB9, 0xBA, 0xBB, 0xBE, 0xBF} and index + 4 < len(data):
            value = struct.unpack_from("<I", data, index + 1)[0]
            if 0 <= value <= MAX_BLOCK_SIZE:
                rows.append({
                    "vaHex": hex32(text_va + index),
                    "value": value,
                    "valueHex": hex32(value),
                    "instruction": (
                        "push imm32"
                        if opcode == 0x68
                        else f"mov {REG_NAMES[opcode - 0xB8]}, imm32"
                    ),
                })
            index += 5
            continue
        if opcode == 0xC7 and index + 6 < len(data) and data[index + 1] in {0x45, 0x44}:
            value = struct.unpack_from("<I", data, index + 3)[0]
            if 0 <= value <= MAX_BLOCK_SIZE:
                rows.append({
                    "vaHex": hex32(text_va + index),
                    "value": value,
                    "valueHex": hex32(value),
                    "instruction": "mov dword ptr [ebp+disp8], imm32",
                })
            index += 7
            continue
        index += 1
    return rows


def call_or_rep_rows(data: bytes, start: int, end: int, text_va: int) -> list[dict]:
    rows = []
    index = start
    while index < min(end, len(data)):
        if data[index] == 0xE8 and index + 4 < len(data):
            rel = struct.unpack_from("<i", data, index + 1)[0]
            target = text_va + index + 5 + rel
            rows.append({
                "vaHex": hex32(text_va + index),
                "kind": "call-rel32",
                "targetHex": hex32(target),
            })
            index += 5
            continue
        if data[index] == 0xFF and index + 1 < len(data) and data[index + 1] == 0x15:
            rows.append({
                "vaHex": hex32(text_va + index),
                "kind": "call-iat",
                "targetHex": None,
            })
            index += 6
            continue
        if data[index] == 0xF3 and index + 1 < len(data) and data[index + 1] in {0xA4, 0xA5, 0xAA, 0xAB}:
            rows.append({
                "vaHex": hex32(text_va + index),
                "kind": f"rep-opcode-0x{data[index + 1]:02x}",
                "targetHex": None,
            })
            index += 2
            continue
        if data[index] in {0xAA, 0xAB}:
            rows.append({
                "vaHex": hex32(text_va + index),
                "kind": f"stos-opcode-0x{data[index]:02x}",
                "targetHex": None,
            })
        index += 1
    return rows


def build_rows(exe: bytes) -> list[dict]:
    text, data = text_section(exe)
    rows = []
    for address in range(BLOCK_SCAN_START, BLOCK_SCAN_END + 1):
        for pos in find_dword_refs(data, address):
            ref = classify_ref(data, pos)
            instruction_index = max(0, pos - (ref.get("instructionPrefixBytes") or 0))
            instruction_end = pos + 4
            required_size = MODE1_SOURCE - address + 1
            direct_covering_write = (
                ref["accessKind"] == "write"
                and isinstance(ref.get("width"), int)
                and address <= MODE1_SOURCE < address + ref["width"]
            )
            address_like = ref["accessKind"] == "address" and address <= MODE1_SOURCE
            window_start = instruction_end
            window_end = instruction_end + WINDOW_BYTES
            size_rows = [
                row for row in small_immediates(data, window_start, window_end, text["va"])
                if row["value"] >= required_size
            ] if address_like else []
            transfer_rows = call_or_rep_rows(data, window_start, window_end, text["va"]) if address_like else []
            block_candidate = address_like and bool(size_rows) and bool(transfer_rows)
            rows.append({
                "addressHex": hex32(address),
                "refVaHex": hex32(text["va"] + pos),
                "instructionVaHex": hex32(text["va"] + instruction_index),
                "accessKind": ref["accessKind"],
                "width": ref["width"],
                "instructionKind": ref["instructionKind"],
                "instruction": ref["instruction"],
                "requiredSizeToCoverMode1": required_size,
                "requiredSizeToCoverMode1Hex": hex32(required_size),
                "directCoveringWrite": direct_covering_write,
                "addressLikeCoveringBase": address_like,
                "coveringSizeCandidates": size_rows,
                "transferCandidates": transfer_rows,
                "blockWriteCandidate": block_candidate,
            })
    rows.sort(key=lambda row: (row["addressHex"], row["refVaHex"]))
    return rows


def build_summary(exe: bytes) -> dict:
    rows = build_rows(exe)
    address_like = [row for row in rows if row["addressLikeCoveringBase"]]
    direct_covering = [row for row in rows if row["directCoveringWrite"]]
    block_candidates = [row for row in rows if row["blockWriteCandidate"]]
    if not direct_covering and not block_candidates:
        address_note = (
            f" It did find {len(address_like)} address-like covering base ref(s), but none has a nearby "
            "covering size immediate, so they are not block-write candidates."
            if address_like
            else " It found no address-like covering base candidate."
        )
        conclusion = (
            "No broad block write candidate was found for 0x0059e348. The scan covers dword immediates in "
            "0x0059e300..0x0059e348 and looks for address-like uses with nearby covering size immediates and "
            "call/rep transfer instructions; it also checks direct wide stores that could cover the mode1 source."
            f"{address_note} This does not replace a runtime trace, but it removes the obvious static "
            "memset/memcpy-style producer shape from the current evidence."
        )
    else:
        conclusion = "At least one broad block write candidate was found; inspect candidate rows before treating 0x0059e348 as unresolved."
    return {
        "scope": "dword references to 0x0059e300..0x0059e348 with nearby size/call/rep block-write scan",
        "mode1SourceHex": hex32(MODE1_SOURCE),
        "scanRangeHex": f"{hex32(BLOCK_SCAN_START)}..{hex32(BLOCK_SCAN_END)}",
        "windowBytes": WINDOW_BYTES,
        "rowCount": len(rows),
        "addressLikeCoveringBaseCount": len(address_like),
        "directCoveringWriteCount": len(direct_covering),
        "blockWriteCandidateCount": len(block_candidates),
        "addressLikeCoveringBases": address_like,
        "directCoveringWrites": direct_covering,
        "blockWriteCandidates": block_candidates,
        "rows": rows,
        "proofFound": False,
        "opcode24Mode1BlockWriteProofFound": False,
        "failedOpcode24Mode1BlockWriteGateIds": FAILED_OPCODE24_MODE1_BLOCK_WRITE_GATE_IDS,
        "missingEvidence": OPCODE24_MODE1_BLOCK_WRITE_MISSING_EVIDENCE,
        "evidenceRefs": OPCODE24_MODE1_BLOCK_WRITE_EVIDENCE_REFS,
        "evidenceRefCount": len(OPCODE24_MODE1_BLOCK_WRITE_EVIDENCE_REFS),
        "promotionStatus": "blocked",
        "conclusion": conclusion,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Opcode 0x24 Mode1 Block Writes",
        "",
        f"Scope: {summary['scope']}.",
        "",
        f"- mode1 source: `{summary['mode1SourceHex']}`",
        f"- scan range: `{summary['scanRangeHex']}`",
        f"- scan window bytes: {summary['windowBytes']}",
        f"- rows: {summary['rowCount']}",
        f"- address-like covering bases: {summary['addressLikeCoveringBaseCount']}",
        f"- direct covering writes: {summary['directCoveringWriteCount']}",
        f"- block write candidates: {summary['blockWriteCandidateCount']}",
        f"- proofFound: `{summary['proofFound']}`",
        f"- opcode24Mode1BlockWriteProofFound: `{summary['opcode24Mode1BlockWriteProofFound']}`",
        f"- promotion status: `{summary['promotionStatus']}`",
        "",
        summary["conclusion"],
        "",
        "## Failed Gates",
        "",
    ]
    lines.extend(f"- `{item}`" for item in summary["failedOpcode24Mode1BlockWriteGateIds"])
    lines.extend([
        "",
        "## Missing Evidence",
        "",
    ])
    lines.extend(f"- {item}" for item in summary["missingEvidence"])
    lines.extend([
        "",
        "## Evidence Refs",
        "",
    ])
    lines.extend(
        f"- `{row['path']}`: {row['description']}"
        for row in summary["evidenceRefs"]
    )
    lines.extend([
        "",
        "## Candidate Rows",
        "",
        "| address | ref | kind | access | required size | sizes | transfers | block candidate | instruction |",
        "| --- | --- | --- | --- | ---: | --- | --- | --- | --- |",
    ])
    candidate_rows = [
        row for row in summary["rows"]
        if row["addressLikeCoveringBase"] or row["directCoveringWrite"] or row["blockWriteCandidate"]
    ]
    for row in candidate_rows:
        sizes = ", ".join(f"{item['valueHex']}@{item['vaHex']}" for item in row["coveringSizeCandidates"]) or "-"
        transfers = ", ".join(f"{item['kind']}@{item['vaHex']}" for item in row["transferCandidates"]) or "-"
        lines.append(
            f"| `{row['addressHex']}` | `{row['instructionVaHex']}` | {row['instructionKind']} | "
            f"{row['accessKind']} | `{row['requiredSizeToCoverMode1Hex']}` | {sizes} | {transfers} | "
            f"{row['blockWriteCandidate']} | `{row['instruction']}` |"
        )
    if not candidate_rows:
        lines.append("| - | - | - | - | - | - | - | - | - |")
    lines.extend([
        "",
        "## All References",
        "",
        "| address | ref | kind | access | width | required size | instruction |",
        "| --- | --- | --- | --- | ---: | ---: | --- |",
    ])
    for row in summary["rows"]:
        lines.append(
            f"| `{row['addressHex']}` | `{row['instructionVaHex']}` | {row['instructionKind']} | "
            f"{row['accessKind']} | {row['width'] if row['width'] is not None else '-'} | "
            f"`{row['requiredSizeToCoverMode1Hex']}` | `{row['instruction']}` |"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    candidate_rows = []
    for row in summary["rows"]:
        if not (row["addressLikeCoveringBase"] or row["directCoveringWrite"] or row["blockWriteCandidate"]):
            continue
        sizes = ", ".join(f"{item['valueHex']}@{item['vaHex']}" for item in row["coveringSizeCandidates"]) or "-"
        transfers = ", ".join(f"{item['kind']}@{item['vaHex']}" for item in row["transferCandidates"]) or "-"
        candidate_rows.append(
            "<tr>"
            f"<td><code>{html.escape(row['addressHex'])}</code></td>"
            f"<td><code>{html.escape(row['instructionVaHex'])}</code></td>"
            f"<td>{html.escape(row['instructionKind'])}</td>"
            f"<td>{html.escape(row['accessKind'])}</td>"
            f"<td><code>{html.escape(row['requiredSizeToCoverMode1Hex'])}</code></td>"
            f"<td>{html.escape(sizes)}</td>"
            f"<td>{html.escape(transfers)}</td>"
            f"<td>{row['blockWriteCandidate']}</td>"
            f"<td><code>{html.escape(row['instruction'])}</code></td>"
            "</tr>"
        )
    all_rows = "\n".join(
        "<tr>"
        f"<td><code>{html.escape(row['addressHex'])}</code></td>"
        f"<td><code>{html.escape(row['instructionVaHex'])}</code></td>"
        f"<td>{html.escape(row['instructionKind'])}</td>"
        f"<td>{html.escape(row['accessKind'])}</td>"
        f"<td>{row['width'] if row['width'] is not None else '-'}</td>"
        f"<td><code>{html.escape(row['requiredSizeToCoverMode1Hex'])}</code></td>"
        f"<td><code>{html.escape(row['instruction'])}</code></td>"
        "</tr>"
        for row in summary["rows"]
    )
    failed_gates = "".join(
        f"<li><code>{html.escape(item)}</code></li>"
        for item in summary["failedOpcode24Mode1BlockWriteGateIds"]
    )
    missing_evidence = "".join(
        f"<li>{html.escape(item)}</li>"
        for item in summary["missingEvidence"]
    )
    evidence_refs = "".join(
        f"<li><code>{html.escape(row['path'])}</code>: {html.escape(row['description'])}</li>"
        for row in summary["evidenceRefs"]
    )
    return "\n".join([
        "<!doctype html><meta charset=\"utf-8\"><title>Save Selector Opcode 0x24 Mode1 Block Writes</title>",
        "<style>body{font-family:system-ui,sans-serif;background:#111;color:#eee;max-width:1180px;margin:24px auto}table{border-collapse:collapse}td,th{border:1px solid #444;padding:6px 8px;vertical-align:top}code{color:#9bd4ff}</style>",
        "<h1>Save Selector Opcode 0x24 Mode1 Block Writes</h1>",
        f"<p>Scope: {html.escape(summary['scope'])}.</p>",
        "<ul>",
        f"<li>mode1 source: <code>{html.escape(summary['mode1SourceHex'])}</code></li>",
        f"<li>scan range: <code>{html.escape(summary['scanRangeHex'])}</code></li>",
        f"<li>scan window bytes: {summary['windowBytes']}</li>",
        f"<li>rows: {summary['rowCount']}</li>",
        f"<li>address-like covering bases: {summary['addressLikeCoveringBaseCount']}</li>",
        f"<li>direct covering writes: {summary['directCoveringWriteCount']}</li>",
        f"<li>block write candidates: {summary['blockWriteCandidateCount']}</li>",
        f"<li>proofFound: <code>{html.escape(str(summary['proofFound']))}</code></li>",
        f"<li>promotion status: <code>{html.escape(summary['promotionStatus'])}</code></li>",
        "</ul>",
        f"<p>{html.escape(summary['conclusion'])}</p>",
        "<h2>Failed Gates</h2>",
        f"<ul>{failed_gates}</ul>",
        "<h2>Missing Evidence</h2>",
        f"<ul>{missing_evidence}</ul>",
        "<h2>Evidence Refs</h2>",
        f"<ul>{evidence_refs}</ul>",
        "<h2>Candidate Rows</h2>",
        "<table><thead><tr><th>address</th><th>ref</th><th>kind</th><th>access</th><th>required size</th><th>sizes</th><th>transfers</th><th>block candidate</th><th>instruction</th></tr></thead><tbody>",
        "\n".join(candidate_rows) or '<tr><td colspan="9">No block write candidates.</td></tr>',
        "</tbody></table>",
        "<h2>All References</h2>",
        "<table><thead><tr><th>address</th><th>ref</th><th>kind</th><th>access</th><th>width</th><th>required size</th><th>instruction</th></tr></thead><tbody>",
        all_rows,
        "</tbody></table>",
    ])


def write_outputs(summary: dict, out_dir: Path = OUT, html_out: Path | None = None) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_opcode24_mode1_block_writes.json").write_text(
        json.dumps(summary, ensure_ascii=False, separators=(",", ":")) + "\n",
        encoding="utf-8",
    )
    if html_out is not None:
        html_out.parent.mkdir(parents=True, exist_ok=True)
        html_out.write_text(html_page(summary), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--exe", type=Path, default=ROOT / "Hwanse2.exe")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    parser.add_argument("--html-out", type=Path, default=None, help="Optional HTML report output path.")
    args = parser.parse_args()
    summary = build_summary(args.exe.read_bytes())
    write_outputs(summary, args.out_dir, args.html_out)
    print(f"wrote opcode24 mode1 block writes -> {args.out_dir / 'save_selector_opcode24_mode1_block_writes.json'}")


if __name__ == "__main__":
    main()
