#!/usr/bin/env python3
"""Summarize file-read candidates for opcode 0x24 mode1 source byte."""
from __future__ import annotations

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

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

from probe_exe_scene_tables import read_sections
from summarize_exe_imports import parse_imports


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
GLOBAL_BUFFER_BASE = 0x0059E310
MODE1_SOURCE = 0x0059E348
MODE1_OFFSET = MODE1_SOURCE - GLOBAL_BUFFER_BASE
GLOBAL_SCAN_START = 0x0059E300
GLOBAL_SCAN_END = 0x0059E350
READFILE_IAT_FALLBACK = 0x005A041C
FAILED_OPCODE24_MODE1_FILE_READ_GATE_IDS = [
    "readfile-mode1-source-producer",
    "runtime-file-io-trace",
    "runtime-watchpoint-trace",
]
OPCODE24_MODE1_FILE_READ_MISSING_EVIDENCE = [
    "ReadFile destination/size covering opcode 0x24 mode1 source 0x0059e348",
    "runtime file-I/O trace proving a save/file read initializes 0x0059e348",
    "runtime watchpoint trace proving the 0x0059e348 producer",
]
OPCODE24_MODE1_FILE_READ_EVIDENCE_REFS = [
    {"path": "Hwanse2.exe", "description": "ReadFile IAT call and destination argument scan"},
    {"path": "out/exe_imports.json", "description": "KERNEL32.ReadFile IAT address"},
    {"path": "out/runtime_trace_feasibility.json", "description": "runtime file-I/O/watchpoint trace blocker"},
]


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


def hex8(value: int) -> str:
    return f"0x{value & 0xFF: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 import_iat_va(imports: dict[str, Any], dll_name: str, function_name: str) -> int | None:
    for dll in imports.get("imports") or []:
        if str(dll.get("dll") or "").lower() != dll_name.lower():
            continue
        for function in dll.get("functions") or []:
            if function.get("name") == function_name and function.get("iatVa") is not None:
                return int(function["iatVa"])
    return None


def find_iat_calls(data: bytes, iat_va: int) -> list[int]:
    needle = b"\xff\x15" + struct.pack("<I", iat_va)
    hits: list[int] = []
    pos = data.find(needle)
    while pos >= 0:
        hits.append(pos)
        pos = data.find(needle, pos + 1)
    return hits


def decode_mov_eax_add_push(data: bytes, index: int, text_va: int) -> dict[str, Any] | None:
    if index + 6 > len(data) or data[index] != 0xB8:
        return None
    base = struct.unpack_from("<I", data, index + 1)[0]
    cursor = index + 5
    offset = 0
    offset_instruction = None
    if cursor + 3 <= len(data) and data[cursor] == 0x83 and data[cursor + 1] == 0xC0:
        offset = struct.unpack_from("b", data, cursor + 2)[0]
        offset_instruction = f"add eax, {hex8(offset)}"
        cursor += 3
    elif cursor + 5 <= len(data) and data[cursor] == 0x05:
        offset = struct.unpack_from("<i", data, cursor + 1)[0]
        offset_instruction = f"add eax, {hex32(offset)}"
        cursor += 5
    elif cursor + 6 <= len(data) and data[cursor] == 0x81 and data[cursor + 1] == 0xC0:
        offset = struct.unpack_from("<i", data, cursor + 2)[0]
        offset_instruction = f"add eax, {hex32(offset)}"
        cursor += 6
    if cursor >= len(data) or data[cursor] != 0x50:
        return None
    target = base + offset
    return {
        "setupIndex": index,
        "setupVaHex": hex32(text_va + index),
        "pushVaHex": hex32(text_va + cursor),
        "target": target,
        "targetHex": hex32(target),
        "baseHex": hex32(base),
        "offset": offset,
        "offsetHex": hex32(offset) if abs(offset) > 0x7F else hex8(offset),
        "setupInstruction": "mov eax, imm32"
        + (f"; {offset_instruction}" if offset_instruction else "")
        + "; push eax",
        "setupByteLength": cursor + 1 - index,
    }


def decode_push_imm32(data: bytes, index: int, text_va: int) -> dict[str, Any] | None:
    if index + 5 > len(data) or data[index] != 0x68:
        return None
    target = struct.unpack_from("<I", data, index + 1)[0]
    return {
        "setupIndex": index,
        "setupVaHex": hex32(text_va + index),
        "pushVaHex": hex32(text_va + index),
        "target": target,
        "targetHex": hex32(target),
        "baseHex": None,
        "offset": target - GLOBAL_BUFFER_BASE,
        "offsetHex": hex32((target - GLOBAL_BUFFER_BASE) & 0xFFFFFFFF),
        "setupInstruction": "push imm32",
        "setupByteLength": 5,
    }


def find_global_destination_setup(data: bytes, start: int, end: int, text_va: int) -> dict[str, Any] | None:
    candidates: list[dict[str, Any]] = []
    for index in range(start, min(end, len(data))):
        decoded = decode_mov_eax_add_push(data, index, text_va)
        if decoded is None:
            decoded = decode_push_imm32(data, index, text_va)
        if decoded is None:
            continue
        target = int(decoded["target"])
        if GLOBAL_SCAN_START <= target <= GLOBAL_SCAN_END:
            candidates.append(decoded)
    if not candidates:
        return None
    return max(candidates, key=lambda row: row["setupIndex"])


def find_read_size_push(data: bytes, start: int, dest_index: int, text_va: int) -> dict[str, Any] | None:
    rows: list[dict[str, Any]] = []
    index = max(start, dest_index - 32)
    while index < dest_index:
        opcode = data[index]
        if opcode == 0x6A and index + 1 < dest_index:
            value = struct.unpack_from("b", data, index + 1)[0]
            if value >= 0:
                rows.append({
                    "size": value,
                    "sizeHex": hex8(value),
                    "sizeVaHex": hex32(text_va + index),
                    "sizeInstruction": "push imm8",
                })
            index += 2
            continue
        if opcode == 0x68 and index + 4 < dest_index:
            value = struct.unpack_from("<I", data, index + 1)[0]
            rows.append({
                "size": value,
                "sizeHex": hex32(value),
                "sizeVaHex": hex32(text_va + index),
                "sizeInstruction": "push imm32",
            })
            index += 5
            continue
        index += 1
    if not rows:
        return None
    return rows[-1]


def scan_readfile_global_rows(exe: bytes, imports: dict[str, Any]) -> tuple[int, list[dict[str, Any]], int]:
    text, data = text_section(exe)
    readfile_iat = import_iat_va(imports, "KERNEL32.dll", "ReadFile") or READFILE_IAT_FALLBACK
    call_positions = find_iat_calls(data, readfile_iat)
    rows: list[dict[str, Any]] = []
    for call_pos in call_positions:
        window_start = max(0, call_pos - 96)
        dest = find_global_destination_setup(data, window_start, call_pos, text["va"])
        if dest is None:
            continue
        size = find_read_size_push(data, window_start, int(dest["setupIndex"]), text["va"])
        read_size = int(size["size"]) if size and isinstance(size.get("size"), int) else None
        target = int(dest["target"])
        covers_mode1 = (
            read_size is not None
            and target <= MODE1_SOURCE < target + max(read_size, 1)
        )
        rows.append({
            "callVaHex": hex32(text["va"] + call_pos),
            "readFileIatVaHex": hex32(readfile_iat),
            "destinationSetupVaHex": dest["setupVaHex"],
            "destinationPushVaHex": dest["pushVaHex"],
            "destinationHex": dest["targetHex"],
            "destinationOffsetFromGlobalBase": target - GLOBAL_BUFFER_BASE,
            "destinationOffsetFromGlobalBaseHex": hex32((target - GLOBAL_BUFFER_BASE) & 0xFFFFFFFF)
            if abs(target - GLOBAL_BUFFER_BASE) > 0x7F
            else hex8(target - GLOBAL_BUFFER_BASE),
            "readSize": read_size,
            "readSizeHex": size.get("sizeHex") if size else None,
            "readSizeVaHex": size.get("sizeVaHex") if size else None,
            "readSizeInstruction": size.get("sizeInstruction") if size else None,
            "coversMode1Source": covers_mode1,
            "setupInstruction": dest["setupInstruction"],
        })
    rows.sort(key=lambda row: row["callVaHex"])
    return readfile_iat, rows, len(call_positions)


def build_summary(exe: bytes, imports: dict[str, Any]) -> dict[str, Any]:
    readfile_iat, rows, readfile_call_count = scan_readfile_global_rows(exe, imports)
    mode1_rows = [row for row in rows if row.get("coversMode1Source")]
    conclusion = (
        "The bounded ReadFile argument scan found global-buffer file reads only for offsets 0x11 and 0x12 "
        "of 0x0059e310. Both reads transfer one byte and therefore do not cover opcode 0x24 mode1 source "
        "0x0059e348 (+0x38). This rules out the obvious file-read producer shape seen in the same global "
        "buffer function, but it is still a static argument-pattern scan rather than a full runtime "
        "watchpoint trace."
        if not mode1_rows
        else "At least one ReadFile destination covers 0x0059e348; inspect the file-read rows before treating mode1 source as unresolved."
    )
    return {
        "mode1SourceHex": hex32(MODE1_SOURCE),
        "globalBufferBaseHex": hex32(GLOBAL_BUFFER_BASE),
        "mode1OffsetHex": hex8(MODE1_OFFSET),
        "globalScanRangeHex": f"{hex32(GLOBAL_SCAN_START)}..{hex32(GLOBAL_SCAN_END)}",
        "readFileIatVaHex": hex32(readfile_iat),
        "readFileCallCount": readfile_call_count,
        "globalDestinationReadFileCount": len(rows),
        "mode1FileReadCandidateCount": len(mode1_rows),
        "mode1FileReadProducerFound": bool(mode1_rows),
        "globalDestinationReadFileRows": rows,
        "mode1FileReadCandidates": mode1_rows,
        "proofFound": False,
        "opcode24Mode1FileReadProofFound": False,
        "failedOpcode24Mode1FileReadGateIds": FAILED_OPCODE24_MODE1_FILE_READ_GATE_IDS,
        "missingEvidence": OPCODE24_MODE1_FILE_READ_MISSING_EVIDENCE,
        "evidenceRefs": OPCODE24_MODE1_FILE_READ_EVIDENCE_REFS,
        "evidenceRefCount": len(OPCODE24_MODE1_FILE_READ_EVIDENCE_REFS),
        "promotionStatus": "blocked",
        "conclusion": conclusion,
    }


def markdown(summary: dict[str, Any]) -> str:
    lines = [
        "# Save Selector Opcode 0x24 Mode1 File Read Context",
        "",
        f"- mode1 source: `{summary['mode1SourceHex']}`",
        f"- global buffer base: `{summary['globalBufferBaseHex']}`",
        f"- mode1 offset from global base: `{summary['mode1OffsetHex']}`",
        f"- global scan range: `{summary['globalScanRangeHex']}`",
        f"- ReadFile IAT: `{summary['readFileIatVaHex']}`",
        f"- ReadFile call count: {summary['readFileCallCount']}",
        f"- global-destination ReadFile rows: {summary['globalDestinationReadFileCount']}",
        f"- mode1 file-read candidates: {summary['mode1FileReadCandidateCount']}",
        f"- proofFound: `{summary['proofFound']}`",
        f"- opcode24Mode1FileReadProofFound: `{summary['opcode24Mode1FileReadProofFound']}`",
        f"- promotion status: `{summary['promotionStatus']}`",
        "",
        summary["conclusion"],
        "",
        "## Failed Gates",
        "",
    ]
    lines.extend(f"- `{item}`" for item in summary["failedOpcode24Mode1FileReadGateIds"])
    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([
        "",
        "## Global ReadFile Destinations",
        "",
        "| call | destination setup | destination | offset | size | covers mode1 | setup |",
        "| --- | --- | --- | --- | ---: | --- | --- |",
    ])
    for row in summary["globalDestinationReadFileRows"]:
        lines.append(
            f"| `{row['callVaHex']}` | `{row['destinationSetupVaHex']}` | "
            f"`{row['destinationHex']}` | `{row['destinationOffsetFromGlobalBaseHex']}` | "
            f"{row['readSize'] if row['readSize'] is not None else '-'} | "
            f"{row['coversMode1Source']} | `{row['setupInstruction']}` |"
        )
    if not summary["globalDestinationReadFileRows"]:
        lines.append("| - | - | - | - | - | - | - |")
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict[str, Any]) -> str:
    rows = "\n".join(
        "<tr>"
        f"<td><code>{html.escape(row['callVaHex'])}</code></td>"
        f"<td><code>{html.escape(row['destinationSetupVaHex'])}</code></td>"
        f"<td><code>{html.escape(row['destinationHex'])}</code></td>"
        f"<td><code>{html.escape(row['destinationOffsetFromGlobalBaseHex'])}</code></td>"
        f"<td>{row['readSize'] if row['readSize'] is not None else '-'}</td>"
        f"<td>{row['coversMode1Source']}</td>"
        f"<td><code>{html.escape(row['setupInstruction'])}</code></td>"
        "</tr>"
        for row in summary["globalDestinationReadFileRows"]
    )
    if not rows:
        rows = "<tr><td colspan='7'>none</td></tr>"
    failed_gates = "".join(
        f"<li><code>{html.escape(item)}</code></li>"
        for item in summary["failedOpcode24Mode1FileReadGateIds"]
    )
    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 f"""<!doctype html>
<meta charset="utf-8">
<title>Save Selector Opcode 0x24 Mode1 File Read Context</title>
<style>
body {{ font-family: sans-serif; margin: 2rem; }}
table {{ border-collapse: collapse; width: 100%; }}
th, td {{ border: 1px solid #ccc; padding: 0.35rem 0.5rem; vertical-align: top; }}
code {{ white-space: nowrap; }}
</style>
<h1>Save Selector Opcode 0x24 Mode1 File Read Context</h1>
<ul>
  <li>mode1 source: <code>{html.escape(summary['mode1SourceHex'])}</code></li>
  <li>global buffer base: <code>{html.escape(summary['globalBufferBaseHex'])}</code></li>
  <li>mode1 offset: <code>{html.escape(summary['mode1OffsetHex'])}</code></li>
  <li>global scan range: <code>{html.escape(summary['globalScanRangeHex'])}</code></li>
  <li>ReadFile IAT: <code>{html.escape(summary['readFileIatVaHex'])}</code></li>
  <li>ReadFile call count: {summary['readFileCallCount']}</li>
  <li>global-destination ReadFile rows: {summary['globalDestinationReadFileCount']}</li>
  <li>mode1 file-read candidates: {summary['mode1FileReadCandidateCount']}</li>
  <li>proofFound: <code>{html.escape(str(summary['proofFound']))}</code></li>
  <li>promotion status: <code>{html.escape(summary['promotionStatus'])}</code></li>
</ul>
<p>{html.escape(summary['conclusion'])}</p>
<h2>Failed Gates</h2>
<ul>{failed_gates}</ul>
<h2>Missing Evidence</h2>
<ul>{missing_evidence}</ul>
<h2>Evidence Refs</h2>
<ul>{evidence_refs}</ul>
<h2>Global ReadFile Destinations</h2>
<table><thead><tr><th>call</th><th>destination setup</th><th>destination</th><th>offset</th><th>size</th><th>covers mode1</th><th>setup</th></tr></thead><tbody>
{rows}
</tbody></table>
"""


def write_outputs(summary: dict[str, Any], 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_file_read_context.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(description=__doc__)
    parser.add_argument("--exe", type=Path, default=ROOT / "Hwanse2.exe")
    parser.add_argument("--imports", type=Path, default=OUT / "exe_imports.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    parser.add_argument(
        "--html-out",
        type=Path,
        default=None,
        help="Optional legacy HTML output path. JSON is the default retained artifact.",
    )
    args = parser.parse_args()
    exe = args.exe.read_bytes()
    imports = (
        json.loads(args.imports.read_text(encoding="utf-8"))
        if args.imports.exists()
        else parse_imports(args.exe)
    )
    summary = build_summary(exe, imports)
    json_out = write_outputs(summary, args.out_dir, args.html_out)
    print(
        "wrote opcode24 mode1 file-read context -> "
        f"{json_out}"
    )


if __name__ == "__main__":
    main()