import sys
from os import path
import parse_bits
from collections import OrderedDict
from functools import partial
from pprint import pprint

class BaseRegister:
    def get(self, n, *, descriptor):
        if type(descriptor.bits) is int:
            return (n >> descriptor.bits) & 1
        else:
            high, low = descriptor.bits
            assert high > low
            mask_length = (high - low) + 1
            mask = (1 << mask_length) - 1
            return (n >> low) & mask

    def get_name(self, n, *, descriptor):
        value = self.get(n, descriptor=descriptor)
        return value, *descriptor.possible_values[value]

def parse_register(register_name):
    base = path.dirname(__file__)

    filename = path.join(base, "bits", register_name.lower() + ".txt")
    l = list(parse_bits.parse_file_fields(filename))
    #return OrderedDict(
    #    (d.field_name, d) for d in parse_bits.aggregate(l)
    #)
    cls = type(register_name, (BaseRegister,), {})
    instance = cls()
    descriptors = list(parse_bits.aggregate(l))
    for descriptor in descriptors:
        setattr(instance, descriptor.field_name, partial(instance.get, descriptor=descriptor))
        setattr(instance, f"_{descriptor.field_name}", partial(instance.get_name, descriptor=descriptor))
        func = getattr(instance, descriptor.field_name)
        for pv_value, (pv_name, _) in descriptor.possible_values.items():
            if pv_name is not None:
                setattr(func, pv_name, pv_value)
    assert getattr(instance, "descriptors", None) is None
    instance.descriptors = descriptors

    return instance

US_CMN_INST = parse_register("US_CMN_INST")
US_ALU_RGB_ADDR = parse_register("US_ALU_RGB_ADDR")
US_ALU_ALPHA_ADDR = parse_register("US_ALU_ALPHA_ADDR")
US_ALU_RGB_INST = parse_register("US_ALU_RGB_INST")
US_ALU_ALPHA_INST = parse_register("US_ALU_ALPHA_INST")
US_ALU_RGBA_INST = parse_register("US_ALU_RGBA_INST")
US_TEX_INST = parse_register("US_TEX_INST")
US_TEX_ADDR = parse_register("US_TEX_ADDR")
US_TEX_ADDR_DXDY = parse_register("US_TEX_ADDR_DXDY")
US_FC_INST = parse_register("US_FC_INST")
US_FC_ADDR = parse_register("US_FC_ADDR")

def disassemble_addr_inner(register_const, address, const, rel):
    assert rel == 0
    if const == register_const.TEMPORARY:
        if address & (1 << 7):
            value = address & 0x7f
            return f"float({value})"
        else:
            return f"temp[{address}]"
    elif const == register_const.CONSTANT:
        return f"const[{address}]"
    else:
        assert False, const

def disassemble_addr(register, addr, prefix):
    addr0 = register.ADDR0(addr)
    addr0_const = register.ADDR0_CONST(addr)
    addr0_rel = register.ADDR0_REL(addr)
    addr1 = register.ADDR1(addr)
    addr1_const = register.ADDR1_CONST(addr)
    addr1_rel = register.ADDR1_REL(addr)
    addr2 = register.ADDR2(addr)
    addr2_const = register.ADDR2_CONST(addr)
    addr2_rel = register.ADDR2_REL(addr)
    _, srcp_op, _ = register._SRCP_OP(addr)

    s0 = disassemble_addr_inner(register.ADDR0_CONST, addr0, addr0_const, addr0_rel)
    s1 = disassemble_addr_inner(register.ADDR1_CONST, addr1, addr1_const, addr1_rel)
    s2 = disassemble_addr_inner(register.ADDR2_CONST, addr2, addr2_const, addr2_rel)
    sp = srcp_op.lower()
    return [
        f"{prefix}{i} = {s}"
        for i, s in [('0', s0), ('1', s1), ('2', s2), ('p', sp)]
    ]

swizzle_strs = ['r', 'g', 'b', 'a', '0', 'h', '1', '_']
sel_strs = ['0', '1', '2', 'p']

def disassemble_rgb_swizzle_sel(us_alu_rgb_inst, us_alu_rgba_inst):
    rgb_sel_a = US_ALU_RGB_INST.RGB_SEL_A(us_alu_rgb_inst)
    red_swiz_a = US_ALU_RGB_INST.RED_SWIZ_A(us_alu_rgb_inst)
    green_swiz_a = US_ALU_RGB_INST.GREEN_SWIZ_A(us_alu_rgb_inst)
    blue_swiz_a = US_ALU_RGB_INST.BLUE_SWIZ_A(us_alu_rgb_inst)
    rgb_mod_a = US_ALU_RGB_INST.RGB_MOD_A(us_alu_rgb_inst)

    rgb_sel_b = US_ALU_RGB_INST.RGB_SEL_B(us_alu_rgb_inst)
    red_swiz_b = US_ALU_RGB_INST.RED_SWIZ_B(us_alu_rgb_inst)
    green_swiz_b = US_ALU_RGB_INST.GREEN_SWIZ_B(us_alu_rgb_inst)
    blue_swiz_b = US_ALU_RGB_INST.BLUE_SWIZ_B(us_alu_rgb_inst)
    rgb_mod_b = US_ALU_RGB_INST.RGB_MOD_B(us_alu_rgb_inst)

    rgb_sel_c = US_ALU_RGBA_INST.RGB_SEL_C(us_alu_rgba_inst)
    red_swiz_c = US_ALU_RGBA_INST.RED_SWIZ_C(us_alu_rgba_inst)
    green_swiz_c = US_ALU_RGBA_INST.GREEN_SWIZ_C(us_alu_rgba_inst)
    blue_swiz_c = US_ALU_RGBA_INST.BLUE_SWIZ_C(us_alu_rgba_inst)
    rgb_mod_c = US_ALU_RGBA_INST.RGB_MOD_C(us_alu_rgba_inst)

    assert rgb_mod_a == 0
    assert rgb_mod_b == 0
    assert rgb_mod_c == 0

    rgb_swiz_a = ''.join(swizzle_strs[n] for n in [red_swiz_a, green_swiz_a, blue_swiz_a])
    rgb_swiz_b = ''.join(swizzle_strs[n] for n in [red_swiz_b, green_swiz_b, blue_swiz_b])
    rgb_swiz_c = ''.join(swizzle_strs[n] for n in [red_swiz_c, green_swiz_c, blue_swiz_c])

    rgb_swiz = [rgb_swiz_a, rgb_swiz_b, rgb_swiz_c]
    rgb_sels = [sel_strs[n] for n in [rgb_sel_a, rgb_sel_b, rgb_sel_c]]

    return ", ".join(f"rgb{sel}.{swiz}" for swiz, sel in zip(rgb_swiz, rgb_sels))

def disassemble_a_swizzle_sel(us_alu_alpha_inst, us_alu_rgba_inst):
    alpha_sel_a = US_ALU_ALPHA_INST.ALPHA_SEL_A(us_alu_alpha_inst)
    alpha_swiz_a = US_ALU_ALPHA_INST.ALPHA_SWIZ_A(us_alu_alpha_inst)
    alpha_mod_a = US_ALU_ALPHA_INST.ALPHA_MOD_A(us_alu_alpha_inst)

    alpha_sel_b = US_ALU_ALPHA_INST.ALPHA_SEL_B(us_alu_alpha_inst)
    alpha_swiz_b = US_ALU_ALPHA_INST.ALPHA_SWIZ_B(us_alu_alpha_inst)
    alpha_mod_b = US_ALU_ALPHA_INST.ALPHA_MOD_B(us_alu_alpha_inst)

    alpha_sel_c = US_ALU_RGBA_INST.ALPHA_SEL_C(us_alu_rgba_inst)
    alpha_swiz_c = US_ALU_RGBA_INST.ALPHA_SWIZ_C(us_alu_rgba_inst)
    alpha_mod_c = US_ALU_RGBA_INST.ALPHA_MOD_C(us_alu_rgba_inst)

    assert alpha_mod_a == 0
    assert alpha_mod_b == 0
    assert alpha_mod_c == 0

    a_swiz = [swizzle_strs[n] for n in [alpha_swiz_a, alpha_swiz_b, alpha_swiz_c]]
    a_sels = [sel_strs[n] for n in [alpha_sel_a, alpha_sel_b, alpha_sel_c]]

    return ", ".join(f"a{sel}.{swiz}" for swiz, sel in zip(a_swiz, a_sels))

def disassemble_alu(ix, code):
    us_cmn_inst = code[ix + 0]
    us_alu_rgb_addr = code[ix + 1]
    us_alu_alpha_addr = code[ix + 2]
    us_alu_rgb_inst = code[ix + 3]
    us_alu_alpha_inst = code[ix + 4]
    us_alu_rgba_inst = code[ix + 5]

    rgb_addr_strs = disassemble_addr(US_ALU_RGB_ADDR, us_alu_rgb_addr, "rgb")
    a_addr_strs = disassemble_addr(US_ALU_ALPHA_ADDR, us_alu_alpha_addr, "a")
    print(", ".join([*rgb_addr_strs, *a_addr_strs]))

    rgb_swizzle_sel = disassemble_rgb_swizzle_sel(us_alu_rgb_inst, us_alu_rgba_inst)
    a_swizzle_sel = disassemble_a_swizzle_sel(us_alu_alpha_inst, us_alu_rgba_inst)

    #omod
    #target
    #alu_wmask

def disassemble(ix, code):
    value = US_CMN_INST.TYPE(code[ix + 0])
    alu_types = {US_CMN_INST.TYPE.US_INST_TYPE_OUT, US_CMN_INST.TYPE.US_INST_TYPE_ALU}
    if value in alu_types:
        disassemble_alu(ix, code)
    else:
        assert False, US_CMN_INST._TYPE(code[ix + 0])

def parse_hex(s):
    assert s.startswith('0x')
    return int(s.removeprefix('0x'), 16)

if __name__ == "__main__":
    filename = sys.argv[1]
    with open(filename) as f:
        buf = f.read()
    code = [parse_hex(c.strip()) for c in buf.split(',') if c.strip()]
    for i in range(len(code) // 6):
        disassemble(i * 6, code)