#! /usr/bin/env python3

import sys

def align(start, end, alignment):
    """Align (<start>, <end>) on multiples of <alignment>."""
    start = (start + alignment - 1) & ~(alignment - 1)
    end &= ~(alignment - 1)
    return start, end

def hexa(data):
    """Readable notation for bytes/bytearray objects"""
    return " ".join(f"{x:02x}" for x in data)

def vardict(vars):
    """Readable notation for variable dictionaries"""
    return " ".join(f"{k}={v:x}" for k, v in vars.items())

#---

class Pattern:
    def __init__(self, string, nibbles):
        self.string = string
        self.nibbles = nibbles

    def __repr__(self):
        return self.string

def make_pattern(string):
    """Compile a pattern of pure hex intermixed with "<var>" references into a
       list of either ints or strings wrapped as a Pattern object."""

    nibbles = []
    i = 0
    while i < len(string):
        if string[i] == "<":
            end = string.index(">", i+1)
            nibbles.append(string[i+1:end])
            i = end+1
        else:
            nibbles.append(int(string[i], 16))
            i += 1

    return Pattern(string, nibbles)

def match_pattern(p, data, vars):
    """Match <pattern> against <data>, unifying with variables <vars>; returns
       either False or extended set of variables."""

    if 2 * len(data) != len(p.nibbles):
        return False

    pos = 0
    vars = vars.copy()

    for pn in p.nibbles:
        dn = (data[pos // 2] >> (0 if pos % 2 else 4)) & 0xf

        if isinstance(pn, str):
            if pn in vars and vars[pn] != dn:
                return False
            vars[pn] = dn
        elif pn != dn:
            return False

        pos += 1

    return vars

def search_pattern(p, data, start, end, *, alignment=1):
    """Returns a list of (<offset>, <vars>) for all occurrences of pattern <p>
       in <alignment>-aligned positions of <data> between <start> and <end>."""

    start, end = align(start, end, alignment)
    n = len(p.nibbles) // 2

    matches = []

    for offset in range(start, end, alignment):
        v = match_pattern(p, data[offset:offset+n], dict())
        if v != False:
            matches.append((offset, v))

    return matches

def eval_pattern(p, vars):
    """Evaluate <p> by substituting variables from <vars>."""

    data = []
    for n in p.nibbles:
        if isinstance(n, str):
            data.append(vars[n])
        else:
            data.append(n)

    return bytes((data[2*i] << 4) | data[2*i+1]
                 for i in range(len(data) // 2))

#---

class Transform:
    def __init__(self, name, patterns):
        self.name = name
        self.insn = [(make_pattern(p), make_pattern(repl))
                     for p, repl in patterns]

def match_transform(tr, base_vars, data, start, end):
    """Matches all instructions from <tr> in <data> between <start> and <end>,
       returning either (False, []) or a pair (<vars>, [<offset>...])
       indicating where instructions matched and with what variables."""

    start, end = align(start, end, 2)
    vars = base_vars.copy()
    offsets = []

    for p, repl in tr.insn:
        n = len(p.nibbles) // 2

        for offset in range(start, end, 2):
            d = data[offset:offset+n]
            v = match_pattern(p, d, vars)
            if v != False:
                offsets.append(offset)
                vars = v
                break
        else:
            return (False, [])

    return (vars, offsets)

def apply_transform(tr, vars, offsets, data):
    """Apply replacements from <tr> on <data> in-place, using the <offsets> and
       <vars> returned by match_transform()."""

    repls = [repl for p, repl in tr.insn]
    replacements = [(offset, eval_pattern(p, vars))
                    for offset, p in zip(offsets, repls)]

    for offset, contents in replacements:
        n = len(contents)
        src = data[offset:offset+n]
        print(f"  Patching at 0x{offset:08x}: {hexa(src)} -> {hexa(contents)}")
        data[offset:offset+n] = contents

#---

def pcrel_compute(pc, disp, *, size=2):
    """Computes the target of a @(disp,pc) argument."""
    return (pc & ~(size - 1)) + 4 + (disp * size)

def pcrel_find(target, data, start, end):
    """Finds all instructions that load <target> with @(disp,pc) in <data>
       between <start> and <end."""

    start, end = align(start, end, 2)
    occurrences = []

    for pc in range(start, end, 2):

        # mov.w @(disp,pc), rn
        movw_atdisppc_rn = make_pattern("9<N><D1><D2>")
        v = match_pattern(movw_atdisppc_rn, data[pc:pc+2], dict())
        if v != False:
            disp = (v["D1"] << 4) | v["D2"]
            if pcrel_compute(pc, disp, size=2) == target:
                occurrences.append((pc, v["N"]))

        # mov.l @(disp,pc), rn
        movl_atdisppc_rn = make_pattern("d<N><D1><D2>")
        v = match_pattern(movl_atdisppc_rn, data[pc:pc+2], dict())
        if v != False:
            disp = (v["D1"] << 4) | v["D2"]
            if pcrel_compute(pc, disp, size=4) == target:
                occurrences.append((pc, v["N"]))

    return occurrences

#---

ML_display_vram = Transform("ML_display_vram", [
    ("e<A>07",    "e<A>0a"),    # Data register is 10
    ("e<B>00",    "e<B>04"),    # Y-Up mode is 4
    ("e<C>04",    "e<C>08"),    # Command register is 8
    ("6<D><B>3",  "e<D>00"),    # Still 0, but rB was changed; use mov #imm, rn
    ("e<E>c0",    "e<E>80"),    # X-address starts at 0x80
])
ML_set_contrast = Transform("ML_set_contrast", [
    ("e<A>06",    "e<A>06"),    # Old contrast register: 6 (detection only)
    ("2<D0><A>0", "0009"),      # Remove: selection of register 6
    ("2<D1><B>0", "0009"),      # Remove: write contrast value
])
ML_get_contrast = Transform("ML_get_contrast", [
    ("e<A>06",    "e<A>06"),    # Old contrast register: 6 (detection only)
    ("2<D0><A>0", "0009"),      # Remove: selection of register 6
    ("60<D1>0",   "e000"),      # Remove: read contrast value; return 0
])

ALL_TRANSFORMS = [
    ML_display_vram,
    ML_set_contrast,
    ML_get_contrast,
]

def main(g1a_input, g1a_output):
    with open(g1a_input, "rb") as fp:
        data = bytearray(fp.read())

    # Find 0xb4000000 and 0xb4010000

    t6k11_addr1 = make_pattern("b4000000")
    refs1 = search_pattern(t6k11_addr1, data, 0, len(data), alignment=4)

    t6k11_addr2 = make_pattern("b4010000")
    refs2 = search_pattern(t6k11_addr2, data, 0, len(data), alignment=4)

    if len(refs1) == 0:
        print(f"no reference to T6K11 address {t6k11_addr1} found")
        return 1

    for offset, _ in refs1:
        print(f"T6K11 reference address {t6k11_addr1} found at 0x{offset:08x}")
        fun_occ = pcrel_find(offset, data, 0, len(data))

        if len(fun_occ) == 0:
            print("  Never loaded")
        for (pc, n) in fun_occ:
            print(f"  0x{pc:08x}: Load into r{n}")
    print()

    # Find function boundaries around, assuming rts marks the end of a function
    # and the start of the next (which is resonable considering the size and
    # simplicity of these functions)

    # Look up to 32 bytes behind, 128 bytes ahead
    LOOKBEHIND  = 32
    LOOKAHEAD   = 128
    RTS = bytes([0x00, 0x0b]) # 000b rts

    functions = []
    for pc, n in fun_occ:
        print(f"Looking for function around 0x{pc:08x}...")
        start, end = -1, -1

        for i in range(0, LOOKBEHIND // 2):
            if data[pc-2*(i+1):pc-2*i] == RTS:
                start = pc-2*(i+1)+4
                print(f"  Starting at 0x{start:08x} after rts")
                break
        else:
            print(f"  No rts found up to {LOOKBEHIND} bytes behind")

        for i in range(0, LOOKAHEAD // 2):
            if data[pc+2*i:pc+2*(i+1)] == RTS:
                end = pc+2*i+4
                print(f"  Stopping at 0x{end:08x} after rts")
                break
        else:
            print(f"  No rts found up to {LOOKAHEAD} bytes ahead")

        if start >= 0 and end >= 0:
            # Find register holding 0xb4000000
            vars = { "D0": n }
            print(f"  0xb4000000 loaded in r{vars['D0']}")

            # Find register holding 0xb4010000
            for offset, _ in refs2:
                occurrences = pcrel_find(offset, data, start, end)
                if len(occurrences) > 0:
                    vars["D1"] = occurrences[0][1]
                    print(f"  0xb4010000 loaded in r{vars['D1']}")
                    break

            functions.append((start, end, vars))

    patched_functions = 0
    ml_display_vram_found = False

    for start, end, base_vars in functions:
        print(f"\nFunction analysis for 0x{start:08x} ... 0x{end:08x}:")

        for tr in ALL_TRANSFORMS:
            vars, offsets = match_transform(tr, base_vars, data, start, end)
            if vars == False:
                continue

            print(f"  Identified {tr.name} with [{vardict(vars)}]:")
            maxlen = max(len(str(p)) for p, repl in tr.insn)
            for offset, (p, repl) in zip(offsets, tr.insn):
                n = len(p.nibbles) // 2
                d = data[offset:offset+n]
                print(f"    Matched {str(p):{maxlen}s} at 0x{offset:08x}:",
                    hexa(d))

            apply_transform(tr, vars, offsets, data)
            patched_functions += 1
            if tr.name == "ML_display_vram":
                ml_display_vram_found = True

    with open(g1a_output, "wb") as fp:
        fp.write(data)

    print()
    print(f"References found to {t6k11_addr1}:", len(refs1))
    print(f"References found to {t6k11_addr2}:", len(refs2))
    print(f"Functions: {len(fun_occ)} hints, {len(functions)} analyzed, "
        f"{patched_functions} patched")
    print(f"ML_display_vram found:", ml_display_vram_found)

    if not ml_display_vram_found:
        print("ML_display_vram not found!")
        return 1
    elif patched_functions == 0:
        print("Nothing patched!")
        return 1
    else:
        print("Success!")
        return 0

if __name__ == "__main__":
    if len(sys.argv) != 3 or "-h" in sys.argv[1:] or "--help" in sys.argv[1:]:
        print(f"usage: {sys.argv[0]} <input g1a> <output g1a>",file=sys.stderr)
        sys.exit(0)
    sys.exit(main(sys.argv[1], sys.argv[2]))