#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import re
import struct
from pathlib import Path


TABLE_BASE_VA = 0x00610398
ROW_STRIDE = 0x6E
MAX_ROW_COUNT = 2048
IMAGE_BASE = 0x00400000


def load_lng_labels(path: Path) -> dict[int, str]:
    labels: dict[int, str] = {}
    for line in path.read_text(encoding="latin1").splitlines():
        match = re.match(r"\s*(\d+)\s+\"(.*)\"\s*$", line)
        if match:
            labels[int(match.group(1))] = match.group(2)
    return labels


def extract_rows(exe_bytes: bytes, labels: dict[int, str]) -> list[dict[str, object]]:
    table_offset = TABLE_BASE_VA - IMAGE_BASE
    rows: list[dict[str, object]] = []
    for row_index in range(MAX_ROW_COUNT):
        row = exe_bytes[
            table_offset + row_index * ROW_STRIDE : table_offset + (row_index + 1) * ROW_STRIDE
        ]
        if len(row) < ROW_STRIDE:
            break
        descriptor_id = struct.unpack_from("<I", row, 0x04)[0]
        if descriptor_id != row_index:
            break
        label_id = struct.unpack_from("<H", row, 0x6A)[0]
        rows.append(
            {
                "row_index": row_index,
                "descriptor_id": descriptor_id,
                "selector_order": struct.unpack_from("<f", row, 0x00)[0],
                "target_mask_bits": row[0x65],
                "label_id": label_id,
                "label": labels.get(label_id, ""),
                "signature_byte_0x63": row[0x63],
                "signature_byte_0x64": row[0x64],
                "signature_hex_0x63_0x6d": row[0x63:0x6E].hex(),
            }
        )
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description="Extract the RT3 EventEffects descriptor table.")
    parser.add_argument("exe", type=Path)
    parser.add_argument("lng", type=Path)
    parser.add_argument("out", type=Path)
    args = parser.parse_args()

    exe_bytes = args.exe.read_bytes()
    labels = load_lng_labels(args.lng)
    artifact = {
        "table_base_va": f"0x{TABLE_BASE_VA:08x}",
        "row_stride_hex": f"0x{ROW_STRIDE:02x}",
        "binary_path_hint": str(args.exe),
        "language_path_hint": str(args.lng),
        "binary_sha256": hashlib.sha256(exe_bytes).hexdigest(),
    }
    artifact["descriptors"] = extract_rows(exe_bytes, labels)
    artifact["descriptor_count"] = len(artifact["descriptors"])
    args.out.write_text(json.dumps(artifact, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()
