#!/usr/bin/env python3
"""Find .text immediates that reference executable virtual addresses."""
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 find_value_refs(data: bytes, sections: list[dict], value: int) -> list[int]:
    text = next(section for section in sections if section["name"] == ".text")
    raw = data[text["raw"] : text["raw"] + text["raw_size"]]
    needle = struct.pack("<I", value)
    refs = []
    search = 0
    while True:
        hit = raw.find(needle, search)
        if hit < 0:
            break
        ref_va = offset_to_va(sections, text["raw"] + hit)
        if ref_va is not None:
            refs.append(ref_va)
        search = hit + 1
    return refs


def find_range_refs(data: bytes, sections: list[dict], start: int, end: int) -> list[tuple[int, int]]:
    text = next(section for section in sections if section["name"] == ".text")
    raw = data[text["raw"] : text["raw"] + text["raw_size"]]
    refs = []
    for index in range(0, len(raw) - 3):
        value = struct.unpack_from("<I", raw, index)[0]
        if start <= value <= end:
            ref_va = offset_to_va(sections, text["raw"] + index)
            if ref_va is not None:
                refs.append((ref_va, value))
    return refs


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)
    args = parser.parse_args()

    data = args.exe.read_bytes()
    sections = read_sections(data)
    for value in args.value:
        refs = find_value_refs(data, sections, value)
        joined = ", ".join(f"0x{ref:08x}" for ref in refs) or "-"
        print(f"0x{value:08x}: {len(refs)} text ref(s): {joined}")

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


if __name__ == "__main__":
    main()