#!/usr/bin/env python3
"""Scan broad block write candidates that could cover secondaryBranchState."""
from __future__ import annotations

import argparse
import html
import json
from pathlib import Path

from summarize_save_selector_opcode24_mode1_block_writes import (
    call_or_rep_rows,
    classify_ref,
    find_dword_refs,
    hex32,
    small_immediates,
    text_section,
)


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
SECONDARY_START = 0x0059E360
SECONDARY_END = 0x0059E36C
SCAN_START = 0x0059E300
SCAN_END = SECONDARY_END - 1
WINDOW_BYTES = 64


def range_hex(start: int, end_exclusive: int) -> str:
    return f"{hex32(start)}..{hex32(end_exclusive - 1)}"


def overlap(start_a: int, end_a: int, start_b: int, end_b: int) -> bool:
    return max(start_a, start_b) < min(end_a, end_b)


def required_sizes(address: int) -> tuple[int, int]:
    if address <= SECONDARY_START:
        return SECONDARY_START - address + 1, SECONDARY_END - address
    return 1, SECONDARY_END - address


def build_rows(exe: bytes) -> list[dict]:
    text, data = text_section(exe)
    rows = []
    for address in range(SCAN_START, 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_touch, required_full_cover = required_sizes(address)
            direct_overlap_write = (
                ref["accessKind"] == "write"
                and isinstance(ref.get("width"), int)
                and overlap(address, address + ref["width"], SECONDARY_START, SECONDARY_END)
            )
            address_like = ref["accessKind"] == "address" and address < SECONDARY_END
            window_start = max(0, instruction_index - WINDOW_BYTES)
            window_end = min(len(data), instruction_end + WINDOW_BYTES)
            size_rows = [
                row for row in small_immediates(data, window_start, window_end, text["va"])
                if row["value"] >= required_touch
            ] if address_like else []
            full_cover_size_rows = [
                row for row in size_rows
                if row["value"] >= required_full_cover
            ]
            transfer_rows = call_or_rep_rows(data, window_start, window_end, text["va"]) if address_like else []
            touch_block_candidate = address_like and bool(size_rows) and bool(transfer_rows)
            full_cover_block_candidate = address_like and bool(full_cover_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"],
                "requiredSizeToTouchSecondary": required_touch,
                "requiredSizeToTouchSecondaryHex": hex32(required_touch),
                "requiredSizeToCoverSecondary": required_full_cover,
                "requiredSizeToCoverSecondaryHex": hex32(required_full_cover),
                "directOverlapWrite": direct_overlap_write,
                "addressLikeTouchingBase": address_like,
                "nearbySizeCandidates": size_rows,
                "nearbyFullCoverSizeCandidates": full_cover_size_rows,
                "nearbyTransferCandidates": transfer_rows,
                "touchBlockWriteCandidate": touch_block_candidate,
                "fullCoverBlockWriteCandidate": full_cover_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["addressLikeTouchingBase"]]
    direct_overlap = [row for row in rows if row["directOverlapWrite"]]
    touch_candidates = [row for row in rows if row["touchBlockWriteCandidate"]]
    full_cover_candidates = [row for row in rows if row["fullCoverBlockWriteCandidate"]]
    if not direct_overlap and not touch_candidates:
        address_note = (
            f" It found {len(address_like)} address-like base ref(s) that could point at or before the table, "
            "but none has a nearby size immediate large enough to touch the table, so they are not block-write candidates."
            if address_like
            else " It found no address-like base refs that could point at or before the table."
        )
        conclusion = (
            "No broad block write candidate was found for secondaryBranchState 0x0059e360..0x0059e36b. "
            "The scan covers dword immediates in 0x0059e300..0x0059e36b and checks address-like uses with "
            "nearby size immediates plus call/rep/stos transfer instructions, as well as direct stores that "
            "overlap the 12-byte table."
            f"{address_note} This removes the obvious static memset/memcpy-style reset shape, but it still does "
            "not prove selector runtime order or a strict map1_01a source hotspot."
        )
    else:
        conclusion = (
            "At least one secondaryBranchState broad block write candidate was found; inspect candidate rows "
            "before treating the secondary reset scope as statically narrowed."
        )
    return {
        "scope": "dword references to 0x0059e300..0x0059e36b with nearby size/call/rep/stos block-write scan",
        "secondaryRangeHex": range_hex(SECONDARY_START, SECONDARY_END),
        "scanRangeHex": range_hex(SCAN_START, SCAN_END + 1),
        "windowBytes": WINDOW_BYTES,
        "rowCount": len(rows),
        "addressLikeTouchingBaseCount": len(address_like),
        "directOverlapWriteCount": len(direct_overlap),
        "touchBlockWriteCandidateCount": len(touch_candidates),
        "fullCoverBlockWriteCandidateCount": len(full_cover_candidates),
        "addressLikeTouchingBases": address_like,
        "directOverlapWrites": direct_overlap,
        "touchBlockWriteCandidates": touch_candidates,
        "fullCoverBlockWriteCandidates": full_cover_candidates,
        "rows": rows,
        "promotionStatus": "blocked",
        "conclusion": conclusion,
    }


def candidate_rows(summary: dict) -> list[dict]:
    return [
        row for row in summary["rows"]
        if (
            row["addressLikeTouchingBase"]
            or row["directOverlapWrite"]
            or row["touchBlockWriteCandidate"]
            or row["fullCoverBlockWriteCandidate"]
        )
    ]


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Secondary Block Writes",
        "",
        f"Scope: {summary['scope']}.",
        "",
        f"- secondaryBranchState: `{summary['secondaryRangeHex']}`",
        f"- scan range: `{summary['scanRangeHex']}`",
        f"- scan window bytes: {summary['windowBytes']}",
        f"- rows: {summary['rowCount']}",
        f"- address-like touching bases: {summary['addressLikeTouchingBaseCount']}",
        f"- direct overlap writes: {summary['directOverlapWriteCount']}",
        f"- touch block write candidates: {summary['touchBlockWriteCandidateCount']}",
        f"- full-cover block write candidates: {summary['fullCoverBlockWriteCandidateCount']}",
        f"- promotion status: `{summary['promotionStatus']}`",
        "",
        summary["conclusion"],
        "",
        "## Candidate Rows",
        "",
        "| address | ref | kind | access | touch size | cover size | sizes | transfers | touch candidate | full cover | instruction |",
        "| --- | --- | --- | --- | ---: | ---: | --- | --- | --- | --- | --- |",
    ]
    rows = candidate_rows(summary)
    if not rows:
        lines.append("| - | - | - | - | - | - | - | - | - | - | - |")
    for row in rows:
        sizes = ", ".join(
            f"{item['valueHex']}@{item['vaHex']}" for item in row["nearbySizeCandidates"]
        ) or "-"
        transfers = ", ".join(
            f"{item['kind']}@{item['vaHex']}" for item in row["nearbyTransferCandidates"]
        ) or "-"
        lines.append(
            f"| `{row['addressHex']}` | `{row['instructionVaHex']}` | {row['instructionKind']} | "
            f"{row['accessKind']} | `{row['requiredSizeToTouchSecondaryHex']}` | "
            f"`{row['requiredSizeToCoverSecondaryHex']}` | {sizes} | {transfers} | "
            f"{row['touchBlockWriteCandidate']} | {row['fullCoverBlockWriteCandidate']} | "
            f"`{row['instruction']}` |"
        )
    lines.extend([
        "",
        "## All References",
        "",
        "| address | ref | kind | access | width | touch size | cover 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['requiredSizeToTouchSecondaryHex']}` | `{row['requiredSizeToCoverSecondaryHex']}` | "
            f"`{row['instruction']}` |"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    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><code>{html.escape(row['requiredSizeToTouchSecondaryHex'])}</code></td>"
        f"<td><code>{html.escape(row['requiredSizeToCoverSecondaryHex'])}</code></td>"
        f"<td>{html.escape(', '.join(item['valueHex'] + '@' + item['vaHex'] for item in row['nearbySizeCandidates']) or '-')}</td>"
        f"<td>{html.escape(', '.join(item['kind'] + '@' + item['vaHex'] for item in row['nearbyTransferCandidates']) or '-')}</td>"
        f"<td>{row['touchBlockWriteCandidate']}</td>"
        f"<td>{row['fullCoverBlockWriteCandidate']}</td>"
        f"<td><code>{html.escape(row['instruction'])}</code></td>"
        "</tr>"
        for row in candidate_rows(summary)
    )
    if not rows:
        rows = "<tr><td colspan=\"11\">No candidate rows</td></tr>"
    return "\n".join([
        "<!doctype html><meta charset=\"utf-8\"><title>Save Selector Secondary Block Writes</title>",
        "<style>body{font-family:system-ui,sans-serif;background:#111;color:#eee;max-width:1200px;margin:24px auto}table{border-collapse:collapse;width:100%;margin:18px 0}td,th{border:1px solid #444;padding:6px 8px;vertical-align:top}code{color:#9bd4ff}</style>",
        "<h1>Save Selector Secondary Block Writes</h1>",
        f"<p>secondaryBranchState <code>{html.escape(summary['secondaryRangeHex'])}</code>; scan <code>{html.escape(summary['scanRangeHex'])}</code>; promotion status <code>{html.escape(summary['promotionStatus'])}</code>.</p>",
        f"<p>rows: {summary['rowCount']}; address-like touching bases: {summary['addressLikeTouchingBaseCount']}; direct overlap writes: {summary['directOverlapWriteCount']}; touch block write candidates: {summary['touchBlockWriteCandidateCount']}; full-cover block write candidates: {summary['fullCoverBlockWriteCandidateCount']}.</p>",
        f"<p>{html.escape(summary['conclusion'])}</p>",
        "<table><thead><tr><th>address</th><th>ref</th><th>kind</th><th>access</th><th>touch size</th><th>cover size</th><th>sizes</th><th>transfers</th><th>touch candidate</th><th>full cover</th><th>instruction</th></tr></thead><tbody>",
        rows,
        "</tbody></table>",
    ])


def write_outputs(summary: dict, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_secondary_block_writes.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (out_dir / "save_selector_secondary_block_writes.html").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)
    args = parser.parse_args()
    summary = build_summary(args.exe.read_bytes())
    write_outputs(summary, args.out_dir)
    print(f"wrote secondary block writes -> {args.out_dir / 'save_selector_secondary_block_writes.html'}")


if __name__ == "__main__":
    main()