#!/usr/bin/env python3
"""Specialized, byte-guarded static lift of one Factorio 2.0.77 function.

This is NOT a general ARM64 lifter or a translated game engine. Original code
and constants are read locally; no game binary is stored in this source file.
"""
import argparse
import hashlib
import json
import struct
from pathlib import Path

FUNCTION_VA = 0x1004DFD14
CONSTANT_VA = 0x10296B4E8
SEED_VA = 0x101AA1848
SEED_SIZE = 24
EXPECTED_SEED = "b1d9c0b86c5590e9c62c7e433ef7b8635d963dd7db98939ebda558004260c7ef"
BODY_SIZE = 104
EXPECTED_BODY = "4fd978c1f02f6c828413991ba37e9d84f5a7805d4c75603ff61f7ef9759db9d3"


def arm64_image(data):
    if data[:4] == bytes.fromhex("cffaedfe"):
        if struct.unpack_from("<I", data, 4)[0] != 0x100000C:
            raise ValueError("Expected ARM64 Mach-O")
        return data, 0
    if data[:4] != bytes.fromhex("cafebabe"):
        raise ValueError("Expected FAT32 or thin ARM64 Mach-O")
    count = struct.unpack_from(">I", data, 4)[0]
    if count > 32 or 8 + count * 20 > len(data):
        raise ValueError("Invalid FAT table")
    for i in range(count):
        cpu, _, offset, size, _ = struct.unpack_from(">IIIII", data, 8 + i * 20)
        if cpu == 0x100000C:
            if offset + size > len(data):
                raise ValueError("Truncated ARM64 slice")
            return data[offset:offset + size], offset
    raise ValueError("No ARM64 slice")


def file_range(image, address, size):
    cursor = 32
    for _ in range(struct.unpack_from("<I", image, 16)[0]):
        command, length = struct.unpack_from("<II", image, cursor)
        if length < 8 or cursor + length > len(image):
            raise ValueError("Invalid Mach-O load command")
        if command == 0x19:
            va, _, offset, file_size = struct.unpack_from("<QQQQ", image, cursor + 24)
            if va <= address and address + size <= va + file_size:
                result = offset + address - va
                if result + size > len(image):
                    raise ValueError("Truncated segment")
                return result
        cursor += length
    raise ValueError("Requested range is not backed by one file segment")


def generate(binary, output):
    data = binary.read_bytes()
    image, slice_offset = arm64_image(data)
    body_offset = file_range(image, FUNCTION_VA, BODY_SIZE)
    body = image[body_offset:body_offset + BODY_SIZE]
    if hashlib.sha256(body).hexdigest() != EXPECTED_BODY:
        raise ValueError("Unrecognized function bytes; refusing to translate another build")
    seed_offset = file_range(image, SEED_VA, SEED_SIZE)
    seed_body = image[seed_offset:seed_offset + SEED_SIZE]
    if hashlib.sha256(seed_body).hexdigest() != EXPECTED_SEED:
        raise ValueError("Unrecognized seed constructor bytes")
    constants_offset = file_range(image, CONSTANT_VA, 32)
    constants = image[constants_offset:constants_offset + 32]
    shifts = struct.unpack("<6i2I", constants)
    if shifts != (13, 2, -19, -25, 12, 4, 0xFFFFE000, 0xFFFFFF80):
        raise ValueError("Unrecognized SIMD shift/mask constants")
    output.mkdir(parents=True, exist_ok=True)
    # Scalarize two NEON lanes; i32 operations preserve modulo-2^32 arithmetic.
    # bfxil w8,w9,#11,#21 replaces the low 21 bits of (c << 17).
    wat = '''(module
  (memory (export "memory") 1)
  (func (export "seed") (param $state i32) (param $seed i32)
    (local $value i32)
    (local.set $value (select (local.get $seed) (i32.const 341)
      (i32.gt_u (local.get $seed) (i32.const 341))))
    (i32.store (local.get $state) (local.get $value))
    (i32.store offset=4 (local.get $state) (local.get $value))
    (i32.store offset=8 (local.get $state) (local.get $value)))
  (func (export "get_int") (param $state i32) (result i32)
    (local $a i32) (local $b i32) (local $c i32)
    (local.set $a (i32.load (local.get $state)))
    (local.set $b (i32.load offset=4 (local.get $state)))
    (local.set $c (i32.load offset=8 (local.get $state)))
    (local.set $a
      (i32.or
        (i32.shr_u (i32.xor (i32.shl (local.get $a) (i32.const 13)) (local.get $a)) (i32.const 19))
        (i32.and (i32.shl (local.get $a) (i32.const 12)) (i32.const -8192))))
    (local.set $b
      (i32.or
        (i32.shr_u (i32.xor (i32.shl (local.get $b) (i32.const 2)) (local.get $b)) (i32.const 25))
        (i32.and (i32.shl (local.get $b) (i32.const 4)) (i32.const -128))))
    (local.set $c
      (i32.or
        (i32.and (i32.shl (local.get $c) (i32.const 17)) (i32.const -2097152))
        (i32.shr_u (i32.xor (local.get $c) (i32.shl (local.get $c) (i32.const 3))) (i32.const 11))))
    (i32.store (local.get $state) (local.get $a))
    (i32.store offset=4 (local.get $state) (local.get $b))
    (i32.store offset=8 (local.get $state) (local.get $c))
    (i32.xor (i32.xor (local.get $a) (local.get $b)) (local.get $c))))
'''
    (output / "factorio_rng.wat").write_text(wat)
    # Native oracle: retain all source instructions except the four ADRP page
    # references, relocated to the same extracted constants in this test image.
    assembly = [".text", ".p2align 2", ".globl _factorio_rng_original", "_factorio_rng_original:"]
    for offset in range(0, len(body), 4):
        word = struct.unpack_from("<I", body, offset)[0]
        if offset in (4, 20, 32, 44):
            assembly.append("  adrp x8, Lrng_constants_page@PAGE")
        else:
            assembly.append(f"  .long 0x{word:08x}")
    assembly += [".p2align 2", ".globl _factorio_rng_seed_original", "_factorio_rng_seed_original:"]
    assembly += [f"  .long 0x{struct.unpack_from('<I', seed_body, i)[0]:08x}" for i in range(0, SEED_SIZE, 4)]
    assembly += [".section __TEXT,__const", ".p2align 12", "Lrng_constants_page:",
                 "  .space 0x4e8", "  .byte " + ",".join(str(x) for x in constants)]
    (output / "oracle.S").write_text("\n".join(assembly) + "\n")
    manifest = {
        "status": "two-function-specialized-lift", "symbol": "RandomGenerator::getInt()",
        "artifact_sha256": hashlib.sha256(data).hexdigest(), "slice_offset": slice_offset,
        "function_va": hex(FUNCTION_VA), "function_size": BODY_SIZE,
        "function_file_offset": slice_offset + body_offset,
        "function_sha256": EXPECTED_BODY, "constants_va": hex(CONSTANT_VA),
        "seed_constructor": {"symbol": "RandomGenerator::RandomGenerator(unsigned int)",
                             "va": hex(SEED_VA), "size": SEED_SIZE,
                             "file_offset": slice_offset + seed_offset, "sha256": EXPECTED_SEED},
        "constants_file_offset": slice_offset + constants_offset,
        "constants_sha256": hashlib.sha256(constants).hexdigest(),
        "translation": "Manually reviewed NEON scalarization, not general automatic lifting",
        "abi": "seed(i32 state_offset, i32 seed); get_int(i32 state_offset) -> i32; 12 state bytes",
        "oracle": "Original instruction words with four ADRP relocations; no game startup or platform calls",
    }
    (output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("binary", type=Path)
    parser.add_argument("--output", type=Path, default=Path("build/rng"))
    args = parser.parse_args()
    generate(args.binary, args.output)