"""DXBC binary dumper — parses container, signatures, SHEX instructions, and validates via D3DDisassemble."""
import struct
import sys
import os

# DXBC opcodes (SM5)
OPCODES = {
    0x00: "add", 0x01: "and", 0x02: "break", 0x03: "breakc",
    0x04: "call", 0x05: "callc", 0x06: "case", 0x07: "continue",
    0x08: "continuec", 0x09: "cut", 0x0A: "default", 0x0B: "deriv_rtx",
    0x0C: "deriv_rty", 0x0D: "discard", 0x0E: "div", 0x0F: "dp2",
    0x10: "dp3", 0x11: "dp4", 0x12: "else", 0x13: "emit",
    0x14: "emitthencut", 0x15: "endif", 0x16: "endloop", 0x17: "endswitch",
    0x18: "eq", 0x19: "exp", 0x1A: "frc", 0x1B: "ftoi",
    0x1C: "ftou", 0x1D: "ge", 0x1E: "iadd", 0x1F: "if",
    0x20: "ieq", 0x21: "ige", 0x22: "ilt", 0x23: "imad",
    0x24: "imax", 0x25: "imin", 0x26: "imul", 0x27: "ine",
    0x28: "ineg", 0x29: "ishl", 0x2A: "ishr", 0x2B: "itof",
    0x2C: "label", 0x2D: "ld", 0x2E: "ld_ms", 0x2F: "log",
    0x30: "loop", 0x31: "lt", 0x32: "mad", 0x33: "min",
    0x34: "max", 0x35: "customdata", 0x36: "mov", 0x37: "movc",
    0x38: "mul", 0x39: "ne", 0x3A: "nop", 0x3B: "not",
    0x3C: "or", 0x3D: "resinfo", 0x3E: "ret", 0x3F: "retc",
    0x40: "round_ne", 0x41: "round_ni", 0x42: "round_pi", 0x43: "round_z",
    0x44: "rsq", 0x45: "sample", 0x46: "sample_c", 0x47: "sample_c_lz",
    0x48: "sample_l", 0x49: "sample_d", 0x4A: "sample_b",
    0x4E: "sqrt", 0x4F: "switch", 0x50: "sincos",
    0x55: "dcl_resource", 0x56: "dcl_constantbuffer",
    0x57: "dcl_sampler", 0x58: "dcl_index_range",
    0x59: "dcl_output_topology", 0x5A: "dcl_input_primitive",
    0x5B: "dcl_max_output_vertex_count",
    0x5C: "dcl_input", 0x5D: "dcl_input_sgv", 0x5E: "dcl_input_siv",
    0x5F: "dcl_input_ps", 0x60: "dcl_input_ps_sgv", 0x61: "dcl_input_ps_siv",
    0x62: "dcl_output", 0x63: "dcl_output_sgv", 0x64: "dcl_output_siv",
    0x65: "dcl_temps", 0x66: "dcl_indexable_temp",
    0x67: "dcl_global_flags",
    0x68: "dcl_thread_group",
    # Extended
    0x6A: "dcl_tgsm_raw", 0x6B: "dcl_tgsm_structured",
    0x6C: "dcl_resource_structured", 0x6D: "dcl_resource_raw",
    0x6E: "dcl_uav_typed", 0x6F: "dcl_uav_raw",
    0x70: "dcl_uav_structured",
    0x7E: "uge", 0x7F: "ult", 0x80: "umul", 0x81: "umad",
    0x82: "umax", 0x83: "umin", 0x84: "ushr",
    0x85: "utof",
    0x8B: "ld_structured",
    0x8C: "store_structured",
    0xA8: "store_uav_typed",
}

OPERAND_TYPES = {
    0: "r", 1: "v", 2: "o", 3: "x(indexable_temp)",
    4: "l(imm32)", 5: "l(imm64)", 6: "s", 7: "t",
    8: "cb", 9: "x(icb)", 10: "x(label)", 11: "x(output_depth)",
    12: "null", 13: "x(rasterizer)", 14: "x(coverage)",
    0x1A: "x(tgsm_stride)",
}

SYS_VALUES = {
    0: "NONE", 1: "POSITION", 2: "CLIP_DISTANCE", 3: "CULL_DISTANCE",
    4: "RENDER_TARGET_ARRAY", 5: "VIEWPORT_ARRAY", 6: "VERTEX_ID",
    7: "PRIMITIVE_ID", 8: "INSTANCE_ID", 9: "IS_FRONT_FACE",
    10: "SAMPLE_INDEX", 11: "FINAL_QUAD_EDGE_TESSFACTOR",
    12: "FINAL_QUAD_INSIDE_TESSFACTOR", 13: "FINAL_TRI_EDGE_TESSFACTOR",
    14: "FINAL_TRI_INSIDE_TESSFACTOR", 15: "FINAL_LINE_DETAIL_TESSFACTOR",
    16: "FINAL_LINE_DENSITY_TESSFACTOR",
    0xFFFF: "DEPTH",
}

COMP = "xyzw"

def mask_str(m):
    return "".join(COMP[i] for i in range(4) if m & (1 << i))

def swizzle_str(s):
    return "".join(COMP[(s >> (i*2)) & 3] for i in range(4))

def decode_operand(data, pos):
    """Decode a DXBC operand, return (text, new_pos)."""
    if pos >= len(data):
        return "???", pos
    token = data[pos]; pos += 1

    num_comp = token & 0x3
    sel_mode = (token >> 2) & 0x3
    op_type = (token >> 12) & 0xFF
    idx_dim = (token >> 20) & 0x3

    # Extended operand?
    extended = (token >> 31) & 1
    negate = False
    abs_val = False
    if extended and pos < len(data):
        ext_token = data[pos]; pos += 1
        negate = bool(ext_token & (1 << 6))
        abs_val = bool(ext_token & (1 << 7))

    type_name = OPERAND_TYPES.get(op_type, f"type{op_type}")

    # Immediate values
    if op_type == 4:  # immediate32
        if num_comp == 1:  # scalar
            if pos < len(data):
                val = struct.unpack_from('<f', struct.pack('<I', data[pos]))[0]
                ival = data[pos]; pos += 1
                return f"l({val})", pos
        elif num_comp == 2:  # vec4
            vals = []
            for _ in range(4):
                if pos < len(data):
                    val = struct.unpack_from('<f', struct.pack('<I', data[pos]))[0]
                    vals.append(f"{val}")
                    pos += 1
            return f"l({','.join(vals)})", pos
        return f"l(?)", pos

    # Component string
    comp_str = ""
    if num_comp == 2:  # vec4
        if sel_mode == 0:  # mask
            m = (token >> 4) & 0xF
            comp_str = "." + mask_str(m)
        elif sel_mode == 1:  # swizzle
            s = (token >> 4) & 0xFF
            sw = swizzle_str(s)
            if len(set(sw)) > 1 or sw != "xyzw":
                comp_str = "." + sw
        elif sel_mode == 2:  # select1
            c = (token >> 4) & 0x3
            comp_str = "." + COMP[c]

    # Index
    indices = []
    for d in range(idx_dim):
        idx_rep = (token >> (22 + d*3)) & 0x7
        if idx_rep == 0:  # immediate32
            if pos < len(data):
                indices.append(str(data[pos])); pos += 1
        elif idx_rep == 1:  # immediate64
            if pos + 1 < len(data):
                val = data[pos] | (data[pos+1] << 32)
                indices.append(str(val)); pos += 2
        elif idx_rep == 2:  # relative
            indices.append("rel")
            _, pos = decode_operand(data, pos)
        elif idx_rep == 3:  # immediate32 + relative
            if pos < len(data):
                base = data[pos]; pos += 1
                indices.append(f"{base}+rel")
                _, pos = decode_operand(data, pos)

    prefix = ""
    suffix = ""
    if negate: prefix = "-"
    if abs_val: prefix = f"|{prefix}"; suffix = "|"

    if op_type == 8:  # cb
        if len(indices) == 2:
            return f"{prefix}cb{indices[0]}[{indices[1]}]{comp_str}{suffix}", pos
        return f"{prefix}cb[{','.join(indices)}]{comp_str}{suffix}", pos

    idx_str = "".join(f"[{i}]" if not i.isdigit() else i for i in indices)
    if len(indices) == 1 and indices[0].isdigit():
        idx_str = indices[0]

    name = type_name.split("(")[0] if "(" not in type_name else type_name
    if op_type in (0, 1, 2, 6, 7):
        return f"{prefix}{type_name}{idx_str}{comp_str}{suffix}", pos
    elif op_type == 12:  # null
        return "null", pos

    return f"{prefix}{type_name}{idx_str}{comp_str}{suffix}", pos


def dump_shex(data, chunk_data):
    """Dump SHEX/SHDR instructions."""
    words = []
    for i in range(0, len(chunk_data), 4):
        words.append(struct.unpack_from('<I', chunk_data, i)[0])

    if len(words) < 2:
        print("  (empty)")
        return

    ver = words[0]
    prog_type = (ver >> 16) & 0xFFFF
    major = (ver >> 4) & 0xF
    minor = ver & 0xF
    types = {0xFFFF: "ps", 0xFFFE: "vs", 0x4353: "cs", 0x4753: "gs", 0x4853: "hs", 0x4453: "ds"}
    pt_str = types.get(prog_type, f'0x{prog_type:04x}')
    print(f"  Version: {pt_str}_{major}_{minor}")
    print(f"  Length: {words[1]} dwords")

    pos = 2
    inst_num = 0
    while pos < len(words):
        opcode_token = words[pos]
        opcode = opcode_token & 0x7FF
        length = (opcode_token >> 24) & 0x7F
        saturate = (opcode_token >> 13) & 1

        if length == 0:
            # Some opcodes have 0 length meaning they're 1 dword
            length = 1

        name = OPCODES.get(opcode, f"op_0x{opcode:02x}")

        # Decode operands
        operand_words = words[pos+1:pos+length]

        sat_str = "_sat" if saturate else ""
        raw_hex = " ".join(f"{w:08x}" for w in words[pos:pos+length])

        # Try to decode operands
        operands = []
        opos = 0
        try:
            while opos < len(operand_words):
                text, new_opos = decode_operand(operand_words, opos)
                if new_opos == opos:
                    break  # no progress
                operands.append(text)
                opos = new_opos
        except:
            pass

        op_str = ", ".join(operands) if operands else ""
        print(f"  {inst_num:4d}: {name}{sat_str} {op_str}")
        print(f"         [{raw_hex}]")

        pos += length
        inst_num += 1


def dump_dxbc(filepath):
    with open(filepath, 'rb') as f:
        data = f.read()

    print(f"File: {filepath} ({len(data)} bytes)")

    if len(data) < 32:
        print("ERROR: File too small for DXBC header")
        return

    magic = data[0:4]
    if magic != b'DXBC':
        print(f"ERROR: Bad magic: {magic}")
        return

    md5 = data[4:20].hex()
    version = struct.unpack_from('<I', data, 20)[0]
    total_size = struct.unpack_from('<I', data, 24)[0]
    chunk_count = struct.unpack_from('<I', data, 28)[0]

    print(f"MD5: {md5}")
    print(f"Version: {version}")
    print(f"Total size: {total_size} (file: {len(data)})")
    print(f"Chunks: {chunk_count}")

    if total_size != len(data):
        print(f"WARNING: Size mismatch!")

    for i in range(chunk_count):
        chunk_offset = struct.unpack_from('<I', data, 32 + i*4)[0]
        chunk_magic = data[chunk_offset:chunk_offset+4].decode('ascii', errors='replace')
        chunk_size = struct.unpack_from('<I', data, chunk_offset+4)[0]
        chunk_data = data[chunk_offset+8:chunk_offset+8+chunk_size]

        print(f"\n{'='*60}")
        print(f"Chunk {i}: {chunk_magic} (offset={chunk_offset}, size={chunk_size})")
        print(f"{'='*60}")

        if chunk_magic in ('ISGN', 'OSGN', 'ISG1', 'OSG1', 'OSG5'):
            elem_count = struct.unpack_from('<I', chunk_data, 0)[0]
            hdr_size = struct.unpack_from('<I', chunk_data, 4)[0]
            print(f"  Elements: {elem_count}, header_size: {hdr_size}")
            for j in range(elem_count):
                base = 8 + j * 24
                name_off = struct.unpack_from('<I', chunk_data, base)[0]
                sem_idx = struct.unpack_from('<I', chunk_data, base+4)[0]
                sys_val = struct.unpack_from('<I', chunk_data, base+8)[0]
                comp_type = struct.unpack_from('<I', chunk_data, base+12)[0]
                reg = struct.unpack_from('<I', chunk_data, base+16)[0]
                mask_rw = struct.unpack_from('<I', chunk_data, base+20)[0]
                mask = mask_rw & 0xFF
                rw = (mask_rw >> 8) & 0xFF

                name_end = name_off
                while name_end < len(chunk_data) and chunk_data[name_end] != 0:
                    name_end += 1
                name = chunk_data[name_off:name_end].decode('utf-8', errors='replace')

                sv_name = SYS_VALUES.get(sys_val, f"?{sys_val}")
                ctypes = {0: "unknown", 1: "uint", 2: "int", 3: "float"}
                ct = ctypes.get(comp_type, f"?{comp_type}")
                io = "v" if chunk_magic == "ISGN" else "o"
                print(f"  [{j}] {name}{sem_idx} : {io}{reg}.{mask_str(mask)} ({ct}) sysval={sv_name} rw={mask_str(rw)}")

                # Validate
                if chunk_magic == "OSGN":
                    for k in range(j):
                        base_k = 8 + k * 24
                        reg_k = struct.unpack_from('<I', chunk_data, base_k+16)[0]
                        if reg_k == reg:
                            mask_k = struct.unpack_from('<I', chunk_data, base_k+20)[0] & 0xFF
                            if mask & mask_k:
                                print(f"    *** ERROR: Register o{reg} collision with element [{k}]!")

        elif chunk_magic in ('SHEX', 'SHDR'):
            dump_shex(data, chunk_data)

        elif chunk_magic == 'RDEF':
            cb_count = struct.unpack_from('<I', chunk_data, 0)[0]
            cb_offset = struct.unpack_from('<I', chunk_data, 4)[0]
            print(f"  Constant buffers: {cb_count}, offset: {cb_offset}")

        elif chunk_magic == 'STAT':
            if len(chunk_data) >= 8:
                inst_count = struct.unpack_from('<I', chunk_data, 0)[0]
                temp_count = struct.unpack_from('<I', chunk_data, 4)[0]
                print(f"  Instructions: {inst_count}, Temp registers: {temp_count}")

    # Try D3DDisassemble
    print(f"\n{'='*60}")
    print("D3DDisassemble validation")
    print(f"{'='*60}")
    try:
        import ctypes as ct
        d3d = ct.windll.d3dcompiler_47
        blob = ct.c_void_p()
        hr = d3d.D3DDisassemble(data, len(data), 0, None, ct.byref(blob))
        if hr == 0:
            vtable_ptr = ct.cast(blob, ct.POINTER(ct.c_void_p))[0]
            vtable = ct.cast(vtable_ptr, ct.POINTER(ct.c_void_p * 5))
            get_ptr = ct.WINFUNCTYPE(ct.c_void_p, ct.c_void_p)(vtable.contents[3])
            get_size = ct.WINFUNCTYPE(ct.c_size_t, ct.c_void_p)(vtable.contents[4])
            ptr = get_ptr(blob)
            size = get_size(blob)
            result = ct.string_at(ptr, size).decode('utf-8')
            print("PASS: D3DDisassemble succeeded")
            print(result)
        else:
            print(f"FAIL: D3DDisassemble returned HRESULT 0x{hr & 0xFFFFFFFF:08X}")
    except Exception as e:
        print(f"SKIP: {e}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <file.dxbc>")
        sys.exit(1)
    dump_dxbc(sys.argv[1])