#!/usr/bin/env python3
"""Scan overlapping stores that may initialize opcode 0x24 mode1 source byte."""
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
SCAN_START = 0x0059E344
SCAN_END = 0x0059E34D
PRODUCER_SCAN_START = 0x0059E340
PRODUCER_SCAN_END = 0x0059E350
FAILED_OPCODE24_MODE1_SOURCE_WRITE_GATE_IDS = [
    "static-mode1-source-write",
    "indexed-mode1-source-producer",
    "runtime-watchpoint-trace",
]
OPCODE24_MODE1_SOURCE_WRITE_MISSING_EVIDENCE = [
    "static overlapping write covering opcode 0x24 mode1 source 0x0059e348",
    "indexed or address-producer candidate that initializes 0x0059e348",
    "runtime watchpoint trace proving the 0x0059e348 producer",
]
OPCODE24_MODE1_SOURCE_WRITE_EVIDENCE_REFS = [
    {"path": "Hwanse2.exe", "description": "absolute-address and static producer-like scans around 0x0059e348"},
    {"path": "out/save_selector_opcode24_mode1_runtime_context.json", "description": "runtime-producer proof gap 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 classify_ref(data: bytes, pos: int, address: int) -> dict:
    candidates = []
    if pos >= 1 and data[pos - 1] in {0xA0, 0xA1, 0xA2, 0xA3}:
        opcode = data[pos - 1]
        if opcode == 0xA0:
            candidates.append(("read", 1, "mov al, ds:[addr]", 1))
        elif opcode == 0xA1:
            candidates.append(("read", 4, "mov eax, ds:[addr]", 1))
        elif opcode == 0xA2:
            candidates.append(("write", 1, "mov ds:[addr], al", 1))
        elif opcode == 0xA3:
            candidates.append(("write", 4, "mov ds:[addr], eax", 1))
    if pos >= 2 and data[pos - 2:pos] in {b"\x8a\x0d", b"\x8b\x0d", b"\x88\x0d", b"\x89\x0d"}:
        prefix = data[pos - 2:pos]
        if prefix == b"\x8a\x0d":
            candidates.append(("read", 1, "mov cl, byte ptr ds:[addr]", 2))
        elif prefix == b"\x8b\x0d":
            candidates.append(("read", 4, "mov ecx, dword ptr ds:[addr]", 2))
        elif prefix == b"\x88\x0d":
            candidates.append(("write", 1, "mov byte ptr ds:[addr], cl", 2))
        elif prefix == b"\x89\x0d":
            candidates.append(("write", 4, "mov dword ptr ds:[addr], ecx", 2))
    if pos >= 2 and data[pos - 2:pos] in {b"\x8a\x15", b"\x8b\x15", b"\x88\x15", b"\x89\x15"}:
        prefix = data[pos - 2:pos]
        if prefix == b"\x8a\x15":
            candidates.append(("read", 1, "mov dl, byte ptr ds:[addr]", 2))
        elif prefix == b"\x8b\x15":
            candidates.append(("read", 4, "mov edx, dword ptr ds:[addr]", 2))
        elif prefix == b"\x88\x15":
            candidates.append(("write", 1, "mov byte ptr ds:[addr], dl", 2))
        elif prefix == b"\x89\x15":
            candidates.append(("write", 4, "mov dword ptr ds:[addr], edx", 2))
    if pos >= 3 and data[pos - 3:pos] == b"\x66\x89\x0d":
        candidates.append(("write", 2, "mov word ptr ds:[addr], cx", 3))
    if pos >= 3 and data[pos - 3:pos] == b"\x66\x8b\x0d":
        candidates.append(("read", 2, "mov cx, word ptr ds:[addr]", 3))
    if pos >= 2 and data[pos - 2:pos] == b"\xc6\x05":
        candidates.append(("write", 1, "mov byte ptr ds:[addr], imm8", 2))
    if pos >= 2 and data[pos - 2:pos] == b"\xc7\x05":
        candidates.append(("write", 4, "mov dword ptr ds:[addr], imm32", 2))
    if not candidates:
        return {
            "accessKind": "unknown",
            "width": None,
            "instruction": "unclassified direct address reference",
            "instructionPrefixBytes": 0,
            "coversMode1Source": address == MODE1_SOURCE,
        }
    access, width, instruction, prefix_len = candidates[0]
    return {
        "accessKind": access,
        "width": width,
        "instruction": instruction,
        "instructionPrefixBytes": prefix_len,
        "coversMode1Source": address <= MODE1_SOURCE < address + width,
    }


def classify_static_producer_candidate(data: bytes, pos: int, address: int) -> dict | None:
    if pos >= 1 and 0xB8 <= data[pos - 1] <= 0xBF:
        reg = ["eax", "ecx", "edx", "ebx", "esp", "ebp", "esi", "edi"][data[pos - 1] - 0xB8]
        return {
            "candidateKind": "address-immediate",
            "accessKind": "address",
            "width": 4,
            "instruction": f"mov {reg}, imm32",
            "instructionPrefixBytes": 1,
            "isWriteCandidate": False,
        }
    if pos >= 1 and data[pos - 1] == 0x68:
        return {
            "candidateKind": "address-immediate",
            "accessKind": "address",
            "width": 4,
            "instruction": "push imm32",
            "instructionPrefixBytes": 1,
            "isWriteCandidate": False,
        }
    if pos >= 2 and data[pos - 2] == 0xC7 and 0xC0 <= data[pos - 1] <= 0xC7:
        reg = ["eax", "ecx", "edx", "ebx", "esp", "ebp", "esi", "edi"][data[pos - 1] - 0xC0]
        return {
            "candidateKind": "address-immediate",
            "accessKind": "address",
            "width": 4,
            "instruction": f"mov {reg}, imm32",
            "instructionPrefixBytes": 2,
            "isWriteCandidate": False,
        }

    if pos >= 2 and data[pos - 2] in {0x88, 0x89, 0x8A, 0x8B, 0xC6, 0xC7}:
        opcode = data[pos - 2]
        modrm = data[pos - 1]
        if 0x80 <= modrm <= 0xBF:
            return {
                "candidateKind": "indexed-displacement",
                "accessKind": "write" if opcode in {0x88, 0x89, 0xC6, 0xC7} else "read",
                "width": 1 if opcode in {0x88, 0x8A, 0xC6} else 4,
                "instruction": f"opcode 0x{opcode:02x} modrm 0x{modrm:02x} disp32",
                "instructionPrefixBytes": 2,
                "isWriteCandidate": opcode in {0x88, 0x89, 0xC6, 0xC7},
            }

    if pos >= 3 and data[pos - 3] in {0x88, 0x89, 0x8A, 0x8B, 0xC6, 0xC7}:
        opcode = data[pos - 3]
        modrm = data[pos - 2]
        sib = data[pos - 1]
        if modrm in {0x04, 0x0C, 0x14, 0x1C, 0x24, 0x2C, 0x34, 0x3C, 0x84, 0x8C, 0x94, 0x9C, 0xA4, 0xAC, 0xB4, 0xBC}:
            return {
                "candidateKind": "sib-displacement",
                "accessKind": "write" if opcode in {0x88, 0x89, 0xC6, 0xC7} else "read",
                "width": 1 if opcode in {0x88, 0x8A, 0xC6} else 4,
                "instruction": f"opcode 0x{opcode:02x} modrm 0x{modrm:02x} sib 0x{sib:02x} disp32",
                "instructionPrefixBytes": 3,
                "isWriteCandidate": opcode in {0x88, 0x89, 0xC6, 0xC7},
            }
    return None


def scan_refs(exe: bytes) -> list[dict]:
    sections = read_sections(exe)
    text = next(section for section in sections if section["name"] == ".text")
    data = exe[text["raw"]: text["raw"] + text["raw_size"]]
    rows = []
    for address in range(SCAN_START, SCAN_END + 1):
        needle = struct.pack("<I", address)
        pos = data.find(needle)
        while pos >= 0:
            ref = classify_ref(data, pos, address)
            rows.append({
                "addressHex": hex32(address),
                "refVaHex": hex32(text["va"] + pos),
                "instructionVaHex": hex32(text["va"] + max(0, pos - (ref.get("instructionPrefixBytes") or 0))),
                "accessKind": ref["accessKind"],
                "width": ref["width"],
                "instruction": ref["instruction"],
                "coversMode1Source": ref["coversMode1Source"],
            })
            pos = data.find(needle, pos + 1)
    rows.sort(key=lambda row: (row["addressHex"], row["refVaHex"]))
    return rows


def scan_static_producer_candidates(exe: bytes) -> list[dict]:
    sections = read_sections(exe)
    text = next(section for section in sections if section["name"] == ".text")
    data = exe[text["raw"]: text["raw"] + text["raw_size"]]
    rows = []
    for address in range(PRODUCER_SCAN_START, PRODUCER_SCAN_END + 1):
        needle = struct.pack("<I", address)
        pos = data.find(needle)
        while pos >= 0:
            candidate = classify_static_producer_candidate(data, pos, address)
            if candidate:
                rows.append({
                    "addressHex": hex32(address),
                    "refVaHex": hex32(text["va"] + pos),
                    "instructionVaHex": hex32(text["va"] + max(0, pos - candidate["instructionPrefixBytes"])),
                    "candidateKind": candidate["candidateKind"],
                    "accessKind": candidate["accessKind"],
                    "width": candidate["width"],
                    "instruction": candidate["instruction"],
                    "isWriteCandidate": candidate["isWriteCandidate"],
                    "exactMode1Displacement": address == MODE1_SOURCE,
                })
            pos = data.find(needle, pos + 1)
    rows.sort(key=lambda row: (row["addressHex"], row["refVaHex"]))
    return rows


def build_summary(exe: bytes) -> dict:
    rows = scan_refs(exe)
    static_candidates = scan_static_producer_candidates(exe)
    covering_writes = [
        row for row in rows
        if row.get("coversMode1Source") and row.get("accessKind") == "write"
    ]
    exact_mode1_refs = [row for row in rows if row.get("addressHex") == hex32(MODE1_SOURCE)]
    indexed_write_candidates = [
        row for row in static_candidates
        if row.get("candidateKind") in {"indexed-displacement", "sib-displacement"}
        and row.get("isWriteCandidate")
    ]
    address_producer_candidates = [
        row for row in static_candidates
        if row.get("candidateKind") == "address-immediate"
    ]
    conclusion = (
        "No overlapping .text store covering 0x0059e348 was found in the absolute-address scan, and no static "
        "address-immediate or indexed-displacement producer candidate was found for the 0x0059e340..0x0059e350 "
        "neighborhood. The only exact mode1 source reference remains the opcode 0x24 mode1 read, so 0x0059e348 "
        "is still an unresolved runtime byte rather than a proven leaf selector. Adjacent references are useful "
        "for object-state context but do not initialize the current mode1 source."
        if not covering_writes and not indexed_write_candidates and not address_producer_candidates
        else "At least one static producer-like candidate was found near 0x0059e348; inspect those rows before treating mode1 source as unresolved."
    )
    return {
        "scope": "absolute .text references to 0x0059e344..0x0059e34d plus static producer-like references to 0x0059e340..0x0059e350",
        "mode1SourceHex": hex32(MODE1_SOURCE),
        "staticProducerScanRangeHex": f"{hex32(PRODUCER_SCAN_START)}..{hex32(PRODUCER_SCAN_END)}",
        "rowCount": len(rows),
        "exactMode1RefCount": len(exact_mode1_refs),
        "coveringWriteCount": len(covering_writes),
        "indexedWriteCandidateCount": len(indexed_write_candidates),
        "addressProducerCandidateCount": len(address_producer_candidates),
        "staticProducerCandidateCount": len(static_candidates),
        "coveringWrites": covering_writes,
        "indexedWriteCandidates": indexed_write_candidates,
        "addressProducerCandidates": address_producer_candidates,
        "staticProducerCandidates": static_candidates,
        "rows": rows,
        "proofFound": False,
        "opcode24Mode1SourceWriteProofFound": False,
        "failedOpcode24Mode1SourceWriteGateIds": FAILED_OPCODE24_MODE1_SOURCE_WRITE_GATE_IDS,
        "missingEvidence": OPCODE24_MODE1_SOURCE_WRITE_MISSING_EVIDENCE,
        "evidenceRefs": OPCODE24_MODE1_SOURCE_WRITE_EVIDENCE_REFS,
        "evidenceRefCount": len(OPCODE24_MODE1_SOURCE_WRITE_EVIDENCE_REFS),
        "promotionStatus": "blocked",
        "conclusion": conclusion,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Opcode 0x24 Mode1 Source Writes",
        "",
        f"Scope: {summary['scope']}.",
        "",
        f"- mode1 source: `{summary['mode1SourceHex']}`",
        f"- static producer scan range: `{summary['staticProducerScanRangeHex']}`",
        f"- exact mode1 refs: {summary['exactMode1RefCount']}",
        f"- overlapping writes covering mode1 source: {summary['coveringWriteCount']}",
        f"- indexed write candidates: {summary['indexedWriteCandidateCount']}",
        f"- address producer candidates: {summary['addressProducerCandidateCount']}",
        f"- proofFound: `{summary['proofFound']}`",
        f"- opcode24Mode1SourceWriteProofFound: `{summary['opcode24Mode1SourceWriteProofFound']}`",
        f"- promotion status: {summary['promotionStatus']}",
        f"- conclusion: {summary['conclusion']}",
        "",
        "## Failed Gates",
        "",
    ]
    lines.extend(f"- `{item}`" for item in summary["failedOpcode24Mode1SourceWriteGateIds"])
    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([
        "",
        "## Static Producer Candidates",
        "",
        "| address | ref | kind | access | width | exact mode1 displacement | instruction |",
        "| --- | --- | --- | --- | ---: | --- | --- |",
    ])
    for row in summary["staticProducerCandidates"]:
        lines.append(
            f"| `{row['addressHex']}` | `{row['instructionVaHex']}` | {row['candidateKind']} | "
            f"{row['accessKind']} | {row['width']} | {row['exactMode1Displacement']} | {row['instruction']} |"
        )
    if not summary["staticProducerCandidates"]:
        lines.append("| - | - | - | - | - | - | - |")
    lines.extend([
        "",
        "## Absolute References",
        "",
        "| address | ref | access | width | covers mode1 | instruction |",
        "| --- | --- | --- | ---: | --- | --- |",
    ])
    for row in summary["rows"]:
        lines.append(
            f"| `{row['addressHex']}` | `{row['instructionVaHex']}` | {row['accessKind']} | "
            f"{row['width'] if row['width'] is not None else '-'} | {row['coversMode1Source']} | {row['instruction']} |"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    producer_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['candidateKind'])}</td>"
        f"<td>{html.escape(row['accessKind'])}</td>"
        f"<td>{row['width']}</td>"
        f"<td>{row['exactMode1Displacement']}</td>"
        f"<td>{html.escape(row['instruction'])}</td>"
        "</tr>"
        for row in summary["staticProducerCandidates"]
    )
    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['accessKind'])}</td>"
        f"<td>{row['width'] if row['width'] is not None else '-'}</td>"
        f"<td>{row['coversMode1Source']}</td>"
        f"<td>{html.escape(row['instruction'])}</td>"
        "</tr>"
        for row in summary["rows"]
    )
    failed_gates = "".join(
        f"<li><code>{html.escape(item)}</code></li>"
        for item in summary["failedOpcode24Mode1SourceWriteGateIds"]
    )
    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 Source Writes</title>",
        "<style>body{font-family:system-ui,sans-serif;background:#111;color:#eee;max-width:1100px;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 Source Writes</h1>",
        f"<p>Scope: {html.escape(summary['scope'])}.</p>",
        f"<p>mode1 source <code>{summary['mode1SourceHex']}</code>; static producer scan <code>{summary['staticProducerScanRangeHex']}</code>; exact refs {summary['exactMode1RefCount']}; overlapping writes covering mode1 source: {summary['coveringWriteCount']}; indexed write candidates: {summary['indexedWriteCandidateCount']}; address producer candidates: {summary['addressProducerCandidateCount']}; proofFound <code>{html.escape(str(summary['proofFound']))}</code>; promotion status: {html.escape(summary['promotionStatus'])}</p>",
        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>Static Producer Candidates</h2>",
        "<table><thead><tr><th>address</th><th>ref</th><th>kind</th><th>access</th><th>width</th><th>exact mode1 displacement</th><th>instruction</th></tr></thead><tbody>",
        producer_rows or '<tr><td colspan="7">No static producer candidates.</td></tr>',
        "</tbody></table>",
        "<h2>Absolute References</h2>",
        "<table><thead><tr><th>address</th><th>ref</th><th>access</th><th>width</th><th>covers mode1</th><th>instruction</th></tr></thead><tbody>",
        rows,
        "</tbody></table>",
    ])


def write_outputs(summary: dict, out_dir: Path, html_out: Path | None = None) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)
    json_out = out_dir / "save_selector_opcode24_mode1_source_writes.json"
    json_out.write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\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")
    return json_out


def main() -> None:
    parser = argparse.ArgumentParser()
    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)
    args = parser.parse_args()
    summary = build_summary(args.exe.read_bytes())
    json_out = write_outputs(summary, args.out_dir, args.html_out)
    print(f"wrote opcode24 mode1 source writes -> {json_out}")


if __name__ == "__main__":
    main()