"""DXBC validation tool - tests blob via D3DDisassemble and supports chunk surgery."""
import struct
import ctypes
import hashlib
import sys

d3d = ctypes.windll.d3dcompiler_47


def test_blob(data, label=""):
    blob = ctypes.c_void_p()
    hr = d3d.D3DDisassemble(bytes(data), len(data), 0, None, ctypes.byref(blob))
    status = "PASS" if hr == 0 else "FAIL"
    print(f"{label}: {status} ({len(data)} bytes)")
    if hr == 0:
        vtable_ptr = ctypes.cast(blob, ctypes.POINTER(ctypes.c_void_p))[0]
        vtable = ctypes.cast(vtable_ptr, ctypes.POINTER(ctypes.c_void_p * 5))
        get_ptr = ctypes.WINFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p)(vtable.contents[3])
        get_size = ctypes.WINFUNCTYPE(ctypes.c_size_t, ctypes.c_void_p)(vtable.contents[4])
        ptr = get_ptr(blob)
        size = get_size(blob)
        return True, ctypes.string_at(ptr, size).decode("utf-8")
    return False, ""


def get_chunks(data):
    count = struct.unpack_from("<I", data, 28)[0]
    chunks = {}
    order = []
    for i in range(count):
        off = struct.unpack_from("<I", data, 32 + i * 4)[0]
        magic = data[off : off + 4].decode("ascii")
        size = struct.unpack_from("<I", data, off + 4)[0]
        chunks[magic] = bytes(data[off + 8 : off + 8 + size])
        order.append(magic)
    return chunks, order


def build_container(chunks_dict, chunk_order):
    num = len(chunk_order)
    header_size = 32 + num * 4
    offsets = []
    cur = header_size
    for name in chunk_order:
        offsets.append(cur)
        cur += 8 + len(chunks_dict[name])
    buf = bytearray(cur)
    # Magic
    buf[0:4] = b"DXBC"
    # MD5 placeholder (bytes 4-19)
    struct.pack_into("<I", buf, 20, 1)  # version
    struct.pack_into("<I", buf, 24, cur)  # total size
    struct.pack_into("<I", buf, 28, num)  # chunk count
    for i, off in enumerate(offsets):
        struct.pack_into("<I", buf, 32 + i * 4, off)
    for (name, off) in zip(chunk_order, offsets):
        buf[off : off + 4] = name.encode("ascii")
        cdata = chunks_dict[name]
        struct.pack_into("<I", buf, off + 4, len(cdata))
        buf[off + 8 : off + 8 + len(cdata)] = cdata
    # Compute MD5 of bytes 20..end
    m = hashlib.md5(bytes(buf[20:]))
    buf[4:20] = m.digest()
    return bytes(buf)


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <file.dxbc> [ref.dxbc]")
        sys.exit(1)

    with open(sys.argv[1], "rb") as f:
        our_data = f.read()

    ok, asm = test_blob(our_data, sys.argv[1])
    if ok:
        print(asm)
        return

    if len(sys.argv) >= 3:
        with open(sys.argv[2], "rb") as f:
            ref_data = f.read()

        ref_c, ref_order = get_chunks(ref_data)
        our_c, our_order = get_chunks(our_data)

        # Verify rebuild works
        rebuilt = build_container(ref_c, ref_order)
        ok, _ = test_blob(rebuilt, "Ref rebuilt (sanity)")

        # Test each chunk replacement
        for chunk in ref_order:
            if chunk in our_c:
                mixed = dict(ref_c)
                mixed[chunk] = our_c[chunk]
                test_blob(build_container(mixed, ref_order), f"Ref + our {chunk}")


if __name__ == "__main__":
    main()