import struct, sys, os
from keystone import Ks, KS_ARCH_X86, KS_MODE_32

ORIG = r"F:\steam\steamapps\common\The I of the Dragon\TheIOfTheDragon.exe"
OUT = r"E:\DragonHD\patch\TheIOfTheDragon_cjk.exe"
TBL = r"E:\DragonHD\patch\tables"

HOOK_DRAW_VA = 0x4EACB0
HOOK_MEAS_VA = 0x4E8026
BACK_DRAW_VA = 0x4EACB8
BACK_MEAS_VA = 0x4E8049

SEC_SIZE = 0x1C0000
OFF_DRAW = 0x0000
OFF_MEAS = 0x0100
OFF_GLYPH = 0x0200
OFF_METRIC = 0x100200


def align(x, a):
    return (x + a - 1) // a * a


def load_tables():
    g = open(os.path.join(TBL, "cjk_glyph.bin"), "rb").read()
    m = open(os.path.join(TBL, "cjk_metric.bin"), "rb").read()
    assert len(g) == 65536 * 16, len(g)
    assert len(m) == 65536 * 8, len(m)
    return g, m


def main():
    data = bytearray(open(ORIG, "rb").read())
    e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
    coff = e_lfanew + 4
    nsec = struct.unpack_from("<H", data, coff + 2)[0]
    optsize = struct.unpack_from("<H", data, coff + 16)[0]
    opt = coff + 20
    imagebase = struct.unpack_from("<I", data, opt + 28)[0]
    sect_align = struct.unpack_from("<I", data, opt + 32)[0]
    file_align = struct.unpack_from("<I", data, opt + 36)[0]
    so = opt + optsize
    secs = []
    for i in range(nsec):
        o = so + i * 40
        vs, va, rs, rp = struct.unpack_from("<IIII", data, o + 8)
        secs.append((va, vs, rp, rs, o))
    last_va, last_vs = secs[-1][0], secs[-1][1]
    new_va = align(last_va + last_vs, sect_align)
    new_raw = align(len(data), file_align)
    sec_abs = imagebase + new_va

    glyph_va = sec_abs + OFF_GLYPH
    metric_va = sec_abs + OFF_METRIC
    draw_va = sec_abs + OFF_DRAW
    meas_va = sec_abs + OFF_MEAS

    ks = Ks(KS_ARCH_X86, KS_MODE_32)
    draw_asm = """
        cmp al, 0x80
        jb ascii
        movzx ecx, al
        shl ecx, 8
        movzx eax, byte ptr [edx + 1]
        or ecx, eax
        mov eax, ecx
        shl eax, 4
        add eax, {glyph}
        mov edx, [eax]
        mov [esi + 0x100C], edx
        mov edx, [eax + 4]
        mov [esi + 0x1010], edx
        mov edx, [eax + 8]
        mov [esi + 0x1014], edx
        mov edx, [eax + 0xC]
        mov [esi + 0x1018], edx
        mov eax, ecx
        shl eax, 3
        add eax, {metric}
        mov edx, [eax]
        mov [esi + 0x1814], edx
        mov edx, [eax + 4]
        mov [esi + 0x1818], edx
        inc dword ptr [esp + 0x10]
        mov ecx, 0xFF
        mov edx, 0xFF
        shl edx, 4
        jmp {back}
    ascii:
        movzx ecx, al
        mov edx, ecx
        shl edx, 4
        jmp {back}
    """.format(glyph=hex(glyph_va), metric=hex(metric_va), back=hex(BACK_DRAW_VA))
    draw_code, _ = ks.asm(draw_asm, draw_va)
    draw_code = bytes(draw_code)

    meas_asm = """
        cmp al, 0x80
        jb ascii
        movzx eax, al
        mov ah, byte ptr [ebx + 1]
        mov ebp, [ecx + 0x14]
        add ebp, dword ptr [eax*8 + {metric}]
        add edx, ebp
        mov ebp, dword ptr [eax*8 + {metric} + 4]
        cmp ebp, esi
        jle skip
        mov esi, ebp
    skip:
        mov ebp, [esp + 0x10]
        inc ebx
        jmp {back}
    ascii:
        mov ebp, [ecx + 0x14]
        movzx eax, al
        add ebp, dword ptr [ecx + eax*8 + 0x101c]
        lea eax, [ecx + eax*8 + 0x101c]
        mov eax, [eax + 4]
        add edx, ebp
        cmp eax, esi
        mov ebp, [esp + 0x10]
        jle skip2
        mov esi, eax
    skip2:
        jmp {back}
    """.format(metric=hex(metric_va), back=hex(BACK_MEAS_VA))
    meas_code, _ = ks.asm(meas_asm, meas_va)
    meas_code = bytes(meas_code)
    print("draw_cave %d bytes @0x%x, metrics_cave %d bytes @0x%x" % (len(draw_code), draw_va, len(meas_code), meas_va))

    if so + (nsec + 1) * 40 > secs[0][2]:
        print("no room for section header", file=sys.stderr)
        sys.exit(3)

    # add section header
    hdr = so + nsec * 40
    struct.pack_into("<8sIIII", data, hdr, b".cjk".ljust(8, b"\x00"), SEC_SIZE, new_va, align(SEC_SIZE, file_align), new_raw)
    struct.pack_into("<IIHH", data, hdr + 36, 0xE0000020, 0, 0, 0)
    struct.pack_into("<H", data, coff + 2, nsec + 1)
    struct.pack_into("<I", data, opt + 56, align(new_va + SEC_SIZE, sect_align))
    if len(data) < new_raw:
        data += b"\x00" * (new_raw - len(data))
    data += b"\x00" * align(SEC_SIZE, file_align)

    sec_off = new_raw
    data[sec_off + OFF_DRAW: sec_off + OFF_DRAW + len(draw_code)] = draw_code
    data[sec_off + OFF_MEAS: sec_off + OFF_MEAS + len(meas_code)] = meas_code

    glyph, metric = load_tables()
    data[sec_off + OFF_GLYPH: sec_off + OFF_GLYPH + len(glyph)] = glyph
    data[sec_off + OFF_METRIC: sec_off + OFF_METRIC + len(metric)] = metric

    def hook(va, target):
        off = va - imagebase
        # .text: raw==va for this exe (.text rp=0x1000, va=0x1000)
        foff = off
        code, _ = ks.asm("jmp %s" % hex(target), va)
        code = bytes(code)
        assert len(code) == 5, code
        data[foff:foff+5] = code
        print("hook @0x%x -> 0x%x" % (va, target))

    hook(HOOK_DRAW_VA, draw_va)
    hook(HOOK_MEAS_VA, meas_va)

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    open(OUT, "wb").write(data)
    print("wrote", OUT, len(data), "bytes")


if __name__ == "__main__":
    main()