import argparse, 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
OFF_LAYOUT_READ = 0x180200
OFF_LAYOUT_APPEND = 0x180300
OFF_LAYOUT_FIT = 0x180400
OFF_TITLE_CENTER = 0x180500

# Fixed left edge at x=620 was tuned for English titles. The panel and Bind
# button are centered at reference x=640 (scaled by the engine from 800px).
TITLE_SITES = ((0x4D748C, 0x4D74A6), (0x4D9A1C, 0x4D9A36),
               (0x4DDAB4, 0x4DDACE), (0x4DE8C8, 0x4DE8E2))

# The wrapped-text path has its own one-byte temporary string and output
# buffer. It must measure and append complete GB2312 characters too.
HOOK_LAYOUT_READ_VA = 0x4ED430
HOOK_LAYOUT_APPEND_VA = 0x4ED50E
HOOK_LAYOUT_FIT_VA = 0x4ED4E9


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


def load_tables(tbl):
    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(orig=ORIG, out=OUT, tbl=TBL):
    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
        shl eax, 8
        mov al, 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))

    # Locals in 0x4ED380: +18..1A = character + NUL, +1B = byte count.
    # +14 is the source byte index, +24 the output byte count. The two
    # formerly unused bytes at +1A/+1B precede the +1C wrap checkpoint.
    # Preserve flags from cmp al,7: the original continuation uses jne.
    layout_read_asm = """
        mov byte ptr [esp + 0x18], al
        mov word ptr [esp + 0x19], 0
        mov byte ptr [esp + 0x1B], 1
        pushfd
        push eax
        cmp dword ptr [esp + 0x2C], 0
        jne checkpoint_ready
        mov eax, ecx
        dec eax
        mov [esp + 0x24], eax
        mov al, byte ptr [esp + 0x20]
    checkpoint_ready:
        cmp al, 0xA1
        jb done
        cmp al, 0xF7
        ja done
        cmp ecx, edx
        jae done
        mov eax, [esp + 0x450]
        mov al, byte ptr [eax + ecx + 1]
        cmp al, 0xA1
        jb done
        cmp al, 0xFE
        ja done
        mov byte ptr [esp + 0x21], al
        mov byte ptr [esp + 0x23], 2
    done:
        pop eax
        popfd
        jmp 0x4ED439
    """
    # Consume the second source byte only AFTER the character fits. On
    # overflow the stock wrap code must retry the lead byte on the next line.
    layout_append_asm = """
        mov eax, [esp + 0x24]
        mov cl, byte ptr [esp + 0x12]
        mov byte ptr [esp + eax + 0x4C], cl
        inc eax
        cmp byte ptr [esp + 0x1B], 2
        jne done
        mov cl, byte ptr [esp + 0x19]
        mov byte ptr [esp + eax + 0x4C], cl
        inc eax
        inc dword ptr [esp + 0x14]
    done:
        mov [esp + 0x24], eax
        mov byte ptr [esp + eax + 0x4C], 0
        mov ecx, [esp + 0x20]
        mov eax, [esp + 0x40]
        cmp eax, ecx
        jle height_done
        mov [esp + 0x20], eax
    height_done:
        mov esi, edx
        cmp byte ptr [esp + 0x1B], 2
        jne appended
        mov eax, [esp + 0x14]
        mov [esp + 0x1C], eax
        mov ebp, esi
        mov eax, [esp + 0x20]
        mov [esp + 0x2C], eax
    appended:
        xor bl, bl
        jmp 0x4ED67C
    """
    # A box narrower than a glyph must still consume the whole character.
    # Emit it on its own line instead of dropping its lead byte. The caller
    # retains its normal clipping behavior for the overwide line.
    layout_fit_asm = """
        lea edx, [eax + esi]
        cmp edx, ecx
        jle fits
        cmp dword ptr [esp + 0x24], 0
        jne overflow
        cmp byte ptr [esp + 0x1B], 2
        jne overflow
    fits:
        jmp 0x4ED4F0
    overflow:
        jmp 0x4ED53D
    """
    layout_read_code = bytes(ks.asm(layout_read_asm, sec_abs + OFF_LAYOUT_READ)[0])
    layout_append_code = bytes(ks.asm(layout_append_asm, sec_abs + OFF_LAYOUT_APPEND)[0])
    layout_fit_code = bytes(ks.asm(layout_fit_asm, sec_abs + OFF_LAYOUT_FIT)[0])
    # Same ABI as DrawText: point/font/color/text/justification. The incoming
    # point is already transformed to screen pixels; subtract the measured
    # width in pixels AFTER transformation (glyphs do not scale with it).
    title_center_asm = """
        pushad
        sub esp, 8
        mov ecx, [esp + 0x30]
        mov eax, [esp + 0x38]
        mov edx, esp
        push 0
        push eax
        push edx
        call 0x4E7FD0
        mov eax, [esp]
        shr eax, 1
        mov edx, [esp + 0x2C]
        sub [edx], eax
        add esp, 8
        popad
        jmp 0x4EAB70
    """
    title_center_code = bytes(ks.asm(title_center_asm, sec_abs + OFF_TITLE_CENTER)[0])
    assert len(draw_code) <= OFF_MEAS - OFF_DRAW
    assert len(meas_code) <= OFF_GLYPH - OFF_MEAS
    assert len(layout_read_code) <= OFF_LAYOUT_APPEND - OFF_LAYOUT_READ
    assert len(layout_append_code) <= OFF_LAYOUT_FIT - OFF_LAYOUT_APPEND
    assert len(layout_fit_code) <= OFF_TITLE_CENTER - OFF_LAYOUT_FIT
    assert OFF_TITLE_CENTER + len(title_center_code) <= SEC_SIZE

    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(tbl)
    data[sec_off + OFF_GLYPH: sec_off + OFF_GLYPH + len(glyph)] = glyph
    data[sec_off + OFF_METRIC: sec_off + OFF_METRIC + len(metric)] = metric
    data[sec_off + OFF_LAYOUT_READ: sec_off + OFF_LAYOUT_READ + len(layout_read_code)] = layout_read_code
    data[sec_off + OFF_LAYOUT_APPEND: sec_off + OFF_LAYOUT_APPEND + len(layout_append_code)] = layout_append_code
    data[sec_off + OFF_LAYOUT_FIT: sec_off + OFF_LAYOUT_FIT + len(layout_fit_code)] = layout_fit_code
    data[sec_off + OFF_TITLE_CENTER: sec_off + OFF_TITLE_CENTER + len(title_center_code)] = title_center_code

    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)
    hook(HOOK_LAYOUT_READ_VA, sec_abs + OFF_LAYOUT_READ)
    hook(HOOK_LAYOUT_APPEND_VA, sec_abs + OFF_LAYOUT_APPEND)
    hook(HOOK_LAYOUT_FIT_VA, sec_abs + OFF_LAYOUT_FIT)
    for anchor, call_site in TITLE_SITES:
        anchor_off, call_off = anchor - imagebase, call_site - imagebase
        expected_call = bytes(ks.asm('call 0x4EAB70', call_site)[0])
        if data[anchor_off:anchor_off + 5] != b'\x68\x6c\x02\x00\x00' or data[call_off:call_off + 5] != expected_call:
            raise ValueError('unexpected option-title instructions at 0x%x' % anchor)
        data[anchor_off:anchor_off + 5] = bytes(ks.asm('push 0x280', anchor)[0])
        data[call_off:call_off + 5] = bytes(ks.asm('call %s' % hex(sec_abs + OFF_TITLE_CENTER), call_site)[0])

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


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--orig", default=ORIG)
    ap.add_argument("--out", default=OUT)
    ap.add_argument("--tbl", default=TBL)
    a = ap.parse_args()
    main(a.orig, a.out, a.tbl)