#!/usr/bin/env python3
"""Find pointers to executable virtual addresses across PE sections."""
from __future__ import annotations

import argparse
import struct
import sys
from pathlib import Path

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

from probe_exe_scene_tables import offset_to_va, read_sections


def parse_int(text: str) -> int:
    return int(text, 0)


def section_for_offset(sections: list[dict], offset: int) -> dict | None:
    for section in sections:
        start = section["raw"]
        end = start + section["raw_size"]
        if start <= offset < end:
            return section
    return None


def find_value_refs(data: bytes, sections: list[dict], value: int, section_names: set[str] | None) -> list[dict]:
    needle = struct.pack("<I", value)
    refs = []
    search = 0
    while True:
        hit = data.find(needle, search)
        if hit < 0:
            break
        search = hit + 1
        section = section_for_offset(sections, hit)
        if section is None:
            continue
        if section_names is not None and section["name"] not in section_names:
            continue
        ref_va = offset_to_va(sections, hit)
        if ref_va is None:
            continue
        refs.append(
            {
                "section": section["name"],
                "fileOffset": hit,
                "refVa": ref_va,
                "value": value,
            }
        )
    return refs


def find_range_refs(
    data: bytes,
    sections: list[dict],
    start: int,
    end: int,
    section_names: set[str] | None,
) -> list[dict]:
    refs = []
    for section in sections:
        if section_names is not None and section["name"] not in section_names:
            continue
        raw_start = section["raw"]
        raw_end = section["raw"] + section["raw_size"]
        raw = data[raw_start:raw_end]
        for index in range(0, len(raw) - 3, 4):
            value = struct.unpack_from("<I", raw, index)[0]
            if not (start <= value <= end):
                continue
            refs.append(
                {
                    "section": section["name"],
                    "fileOffset": raw_start + index,
                    "refVa": section["va"] + index,
                    "value": value,
                }
            )
    return refs


def print_ref(ref: dict) -> None:
    print(
        f"  {ref['section']:7} va=0x{ref['refVa']:08x} "
        f"file=0x{ref['fileOffset']:06x} -> 0x{ref['value']:08x}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--exe", type=Path, default=Path("Hwanse2.exe"))
    parser.add_argument("--value", type=parse_int, action="append", default=[])
    parser.add_argument("--range", nargs=2, metavar=("START", "END"), type=parse_int)
    parser.add_argument(
        "--sections",
        default=".text,.rdata,.data",
        help="comma-separated section names to scan, or 'all'",
    )
    parser.add_argument("--limit", type=int, default=200)
    args = parser.parse_args()

    data = args.exe.read_bytes()
    sections = read_sections(data)
    section_names = None if args.sections == "all" else set(args.sections.split(","))

    for value in args.value:
        refs = find_value_refs(data, sections, value, section_names)
        print(f"0x{value:08x}: {len(refs)} ref(s)")
        for ref in refs[: args.limit]:
            print_ref(ref)
        if len(refs) > args.limit:
            print(f"  ... {len(refs) - args.limit} more")

    if args.range:
        refs = find_range_refs(data, sections, args.range[0], args.range[1], section_names)
        print(f"range 0x{args.range[0]:08x}..0x{args.range[1]:08x}: {len(refs)} ref(s)")
        for ref in refs[: args.limit]:
            print_ref(ref)
        if len(refs) > args.limit:
            print(f"  ... {len(refs) - args.limit} more")


if __name__ == "__main__":
    main()