from __future__ import annotations

import argparse
import csv
import uuid
from pathlib import Path

import pefile


IMAGE_BASE = 0x10000000
PROXY_VTBL_LIST = 0x10007CC8
STUB_VTBL_LIST = 0x10007CE4
NAME_LIST = 0x10007D00
INTERFACE_COUNT = 6


def u16(data: bytes, offset: int) -> int:
    return int.from_bytes(data[offset:offset + 2], "little")


def u32(data: bytes, offset: int) -> int:
    return int.from_bytes(data[offset:offset + 4], "little")


class PeView:
    def __init__(self, path: Path):
        self.pe = pefile.PE(str(path))
        self.data = bytes(self.pe.__data__)
        self.base = self.pe.OPTIONAL_HEADER.ImageBase

    def offset(self, va: int) -> int:
        return self.pe.get_offset_from_rva(va - self.base)

    def u16(self, va: int) -> int:
        return u16(self.data, self.offset(va))

    def u32(self, va: int) -> int:
        return u32(self.data, self.offset(va))

    def guid(self, va: int) -> str:
        return str(uuid.UUID(bytes_le=self.data[self.offset(va):self.offset(va) + 16])).upper()

    def asciiz(self, va: int) -> str:
        start = self.offset(va)
        end = self.data.index(b"\x00", start)
        return self.data[start:end].decode("ascii")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dll",
        type=Path,
        default=Path(r"C:\Program Files (x86)\ArchestrA\Framework\Bin\NmxSvcps.dll"),
    )
    parser.add_argument("--out", type=Path, default=Path("analysis/proxy/nmxsvcps-proxy-layout.tsv"))
    args = parser.parse_args()

    pe = PeView(args.dll)
    args.out.parent.mkdir(parents=True, exist_ok=True)

    with args.out.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle, delimiter="\t", lineterminator="\n")
        writer.writerow([
            "interface_index",
            "name",
            "iid",
            "iid_va",
            "method_count",
            "user_method_count",
            "proxy_vtbl_va",
            "stub_vtbl_va",
            "proxy_info_va",
            "proxy_info_words",
        ])

        for index in range(INTERFACE_COUNT):
            name_va = pe.u32(NAME_LIST + index * 4)
            name = pe.asciiz(name_va)
            proxy_vtbl = pe.u32(PROXY_VTBL_LIST + index * 4)
            stub_vtbl = pe.u32(STUB_VTBL_LIST + index * 4)
            proxy_info = pe.u32(proxy_vtbl)
            iid_va = pe.u32(proxy_vtbl + 4)
            method_count = pe.u32(stub_vtbl + 8)
            user_method_count = max(0, method_count - 3)
            proxy_info_words = [pe.u32(proxy_info + i * 4) for i in range(8)]

            writer.writerow([
                index,
                name,
                pe.guid(iid_va),
                f"0x{iid_va:08x}",
                method_count,
                user_method_count,
                f"0x{proxy_vtbl:08x}",
                f"0x{stub_vtbl:08x}",
                f"0x{proxy_info:08x}",
                ",".join(f"0x{word:08x}" for word in proxy_info_words),
            ])

    print(f"wrote {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())