#!/usr/bin/env python3
"""Напряжение батареи вместо штатной иконки + уборка полосы у всплывающего окна.

Патчи:
  1. Свой код в свободном месте после конца образа (0x08029000):
       stub        — пустышка `bx lr`
       batt_draw   — рисует "8.4V" мелким шрифтом 5x8 на странице 0
       tick_router — вызывает оригинал, затем batt_draw (только на главном экране)
  2. `bl 0x08008300` @0x080148E2 — отрисовка штатной шкалы -> stub (иконки больше нет).
  3. `bl 0x08010A00` @0x0801E2A4 — «напряжение -> индекс 0..3» убрано (2 x nop),
     чтобы флаг перерисовки взводился на каждые 0.1 В, а не только на порогах.
  4. `bl 0x0800842C` @0x080085DE — верхняя полоса рамки всплывающего окна убрана.
  5. `bl 0x0801DFD0` @0x08020806 — последний вызов цепочки отрисовки в ui_tick_normal
     перенаправлен на tick_router.

Почему через конец цикла: раньше мы рисовали из середины кадра (из цепочки обновления
батареи), и любой более поздний вывод затирал цифры — при входящем вызове «7» теряла
перекладину и читалась как «1». Теперь наш текст ложится последним и перекрывает всё.
"""
import struct
from keystone import Ks, KS_ARCH_ARM, KS_MODE_THUMB

BASE = 0x08002800
SRC  = "stock-fw/rt4d_ru_full.bin"
OUT  = "stock-fw/rt4d_ru_batt.bin"

SMALL_CHAR  = 0x080089AC     # void small_char(u8 page, u8 x, u8 ch, u8 dim)
BATT_VOLT   = 0x200008B0     # u16, напряжение * 10
G_SCREEN    = 0x200008B3     # u8, текущий экран (0 = главный)

GAUGE_CALL  = 0x080148E2     # bl 0x08008300 — штатная шкала батареи
GAUGE_FUNC  = 0x08008300
LEVEL_CALL  = 0x0801E2A4     # bl 0x08010A00 — напряжение -> индекс
LEVEL_FUNC  = 0x08010A00
TOPBAR_CALL = 0x080085DE     # bl fill_col_run(page=1) — верхняя полоса рамки
FILL_RUN    = 0x0800842C
TICK_CALL   = 0x08020806     # bl 0x0801DFD0 — последний вызов отрисовки за кадр
TICK_FUNC   = 0x0801DFD0

MYCODE = 0x08029000
X0, X1, X2, X3 = 108, 113, 118, 123     # позиции четырёх символов


def thumb_bl(src: int, dst: int) -> bytes:
    off = dst - (src + 4)
    assert -(1 << 24) <= off < (1 << 24), "цель BL слишком далеко"
    off >>= 1
    s = (off >> 23) & 1
    i1 = (off >> 22) & 1
    i2 = (off >> 21) & 1
    imm10 = (off >> 11) & 0x3FF
    imm11 = off & 0x7FF
    j1 = (~i1 ^ s) & 1
    j2 = (~i2 ^ s) & 1
    return struct.pack("<HH", 0xF000 | (s << 10) | imm10,
                              0xD000 | (j1 << 13) | (j2 << 11) | imm11)


def main():
    img = bytearray(open(SRC, "rb").read())
    print(f"исходный образ: {len(img)} Б, конец на 0x{BASE + len(img) - 1:08X}")
    assert BASE + len(img) == MYCODE, "образ не заканчивается там, где ждём код"
    ks = Ks(KS_ARCH_ARM, KS_MODE_THUMB)

    def check(addr, target, what):
        o = addr - BASE
        cur = bytes(img[o:o + 4])
        assert cur == thumb_bl(addr, target), \
            f"по 0x{addr:08X} ({what}) не тот вызов: {cur.hex(' ')}"
        return o

    o_gauge = check(GAUGE_CALL, GAUGE_FUNC, "шкала батареи")
    o_level = check(LEVEL_CALL, LEVEL_FUNC, "напряжение->индекс")
    o_bar   = check(TOPBAR_CALL, FILL_RUN, "верхняя полоса рамки")
    o_tick  = check(TICK_CALL, TICK_FUNC, "хвост отрисовки")
    print("все четыре точки патча подтверждены")

    # ---- раскладка своего кода ----
    stub_addr = MYCODE
    stub = bytes(ks.asm("bx lr", stub_addr)[0])

    batt_addr = stub_addr + len(stub)
    batt_asm = f"""
        push  {{r4, r5, r6, lr}}
        movw  r4, #{BATT_VOLT & 0xFFFF}
        movt  r4, #{(BATT_VOLT >> 16) & 0xFFFF}
        ldrh  r4, [r4]
        uxtb  r4, r4
        movs  r1, #205
        muls  r1, r4, r1
        lsrs  r5, r1, #11
        movs  r1, #10
        muls  r1, r5, r1
        subs  r4, r4, r1
        movs  r0, #0
        movs  r1, #{X0}
        mov   r2, r5
        add   r2, #0x30
        movs  r3, #0
        bl    #{SMALL_CHAR}
        movs  r0, #0
        movs  r1, #{X1}
        movs  r2, #0x2E
        movs  r3, #0
        bl    #{SMALL_CHAR}
        movs  r0, #0
        movs  r1, #{X2}
        mov   r2, r4
        add   r2, #0x30
        movs  r3, #0
        bl    #{SMALL_CHAR}
        movs  r0, #0
        movs  r1, #{X3}
        movs  r2, #0x56
        movs  r3, #0
        bl    #{SMALL_CHAR}
        pop   {{r4, r5, r6, pc}}
    """
    batt = bytes(ks.asm(batt_asm, batt_addr)[0])

    # tick_router: оригинал -> проверка экрана -> batt_draw.
    # keystone не умеет метки, поэтому условный переход кодируем вручную.
    tick_addr = batt_addr + len(batt)
    partA = bytes(ks.asm(f"""
        push  {{r4, lr}}
        bl    #{TICK_FUNC}
        movw  r4, #{G_SCREEN & 0xFFFF}
        movt  r4, #{(G_SCREEN >> 16) & 0xFFFF}
        ldrb  r4, [r4]
        cmp   r4, #0
    """, tick_addr)[0])
    bne_addr = tick_addr + len(partA)
    partB = bytes(ks.asm(f"bl #{batt_addr}", bne_addr + 2)[0])
    bne = struct.pack("<H", 0xD100 | (((len(partB) - 2) // 2) & 0xFF))   # bne -> за partB
    partC = bytes(ks.asm("pop {r4, pc}", bne_addr + 2 + len(partB))[0])
    tick = partA + bne + partB + partC

    code = stub + batt + tick
    print(f"свой код: {len(code)} Б  (stub 0x{stub_addr:08X}, "
          f"batt_draw 0x{batt_addr:08X}, tick_router 0x{tick_addr:08X})")

    img += code
    while len(img) % 4:
        img.append(0xFF)

    # ---- патчи ----
    img[o_gauge:o_gauge + 4] = thumb_bl(GAUGE_CALL, stub_addr)
    print("штатная шкала батареи отключена")
    img[o_level:o_level + 4] = b"\x00\xbf\x00\xbf"
    print("преобразование в индекс убрано — перерисовка на каждые 0.1 В")
    img[o_bar:o_bar + 4] = b"\x00\xbf\x00\xbf"
    print("верхняя полоса рамки убрана")
    img[o_tick:o_tick + 4] = thumb_bl(TICK_CALL, tick_addr)
    print("вывод напряжения перевешен на конец кадра")

    open(OUT, "wb").write(bytes(img))
    print(f"готово: {OUT}, {len(img)} Б")


if __name__ == "__main__":
    main()