#!/usr/bin/env python3
"""Интерактивный агент в прошивке: управление FM100B через COM-порт с ПК, без перешивки.

Протокол (ПК -> рация по USART6):
    0xAA <len> <len байт «сырого» FM100B-кадра>
агент блитит эти len байт в USART3 (модулю). Ответы модуля идут обратно в порт
через снифер-хук on_frame (кадр: F1 F2 <len> <кадр>). Так я с ПК шлю любой кадр
и сразу вижу реакцию — перебор команд (AMBE, record и пр.) вживую.

Приём с порта — по байту за тик в главном цикле (не блокирует UI, в отличие от зонда).
CPS одновременно не подключать. Основа rt4d_ru_batt.bin, откат — она же.

    python3 build_agent.py
    python3 tools/agent.py send 25 01 01           # спросить версию модуля
    python3 tools/agent.py raw 68 25 01 01 00 00 00 00 10
"""
import struct
from keystone import Ks, KS_ARCH_ARM, KS_MODE_THUMB

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

ON_FRAME     = 0x08006348
FRAME_CALL   = 0x08018DD8
UI_HOOK      = 0x08020806
TICK_FUNC    = 0x0801DFD0
TX_BYTE      = 0x08018BFC     # usart6_tx_byte(u8)
USART3_TXBUF = 0x08006C9C     # usart3_tx_buf(buf,len)
# приём CPS идёт через USART6-ISR, который складывает байты в это кольцо.
# Читаем кольцо по своему указателю — порт/ISR не трогаем, CPS не ломаем.
RING_BUF   = 0x200082EF       # 4K кольцо RX (маска 0xFFF)
RING_HEAD  = 0x20000C60       # u32, растущий индекс записи (из ISR)

MAGIC_ADDR = 0x2000B000
MAGIC_VAL  = 0x41474E54       # 'AGNT'
ST_STATE   = 0x2000B040       # 0=ждём AA, 1=ждём len, 2=собираем
ST_LEN     = 0x2000B041
ST_IDX     = 0x2000B042
ST_BUF     = 0x2000B044       # до 64 байт -> 0x2000B084
MY_TAIL    = 0x2000B088       # u32, наш указатель чтения кольца

EQ, NE, HS, LS, LO = 0, 1, 2, 9, 3


def thumb_bl(src, dst):
    off = (dst - (src + 4)) >> 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)


class Asm:
    def __init__(self, ks, base): self.ks=ks; self.base=base; self.items=[]
    def asm(self,c): self.items.append(("a",c))
    def label(self,n): self.items.append(("l",n))
    def b(self,t,cond=None): self.items.append(("b",cond,t))
    def _sz(self,it):
        if it[0]=="a": return len(bytes(self.ks.asm(it[1],0)[0]))
        if it[0]=="b": return 2
        return 0
    def build(self):
        pos={}; off=0
        for it in self.items:
            if it[0]=="l": pos[it[1]]=off
            off+=self._sz(it)
        out=bytearray(); off=0
        for it in self.items:
            if it[0]=="a": out+=bytes(self.ks.asm(it[1],self.base+off)[0])
            elif it[0]=="b":
                _,cond,tgt=it; d=(self.base+pos[tgt])-(self.base+off+4)
                out+= struct.pack("<H",0xE000|((d>>1)&0x7FF)) if cond is None \
                      else struct.pack("<H",0xD000|(cond<<8)|((d>>1)&0xFF))
            off+=self._sz(it)
        return bytes(out)


def main():
    img=bytearray(open(SRC,"rb").read())
    end=BASE+len(img); code_at=(end+3)&~3
    img+=b"\xff"*(code_at-end)
    ks=Ks(KS_ARCH_ARM,KS_MODE_THUMB)
    def lo(a): return a&0xFFFF
    def hi(a): return (a>>16)&0xFFFF

    # ---- agent_poll: читать кольцо RX по своему указателю ----
    # r7 = my_tail (u32). Пока my_tail != head — брать ring[my_tail & 0xFFF].
    poll_at=code_at
    a=Asm(ks,poll_at)
    a.asm(f"""
        push {{r4, r5, r6, r7, lr}}
        movw r7, #{lo(MY_TAIL)}
        movt r7, #{hi(MY_TAIL)}
        ldr r7, [r7]
    """)
    a.label("ploop")
    # head?
    a.asm(f"""
        movw r0, #{lo(RING_HEAD)}
        movt r0, #{hi(RING_HEAD)}
        ldr r0, [r0]
        cmp r7, r0
    """)
    a.b("pret", EQ)
    # byte = ring[my_tail & 0xFFF]; my_tail++
    a.asm(f"""
        movw r0, #{lo(RING_BUF)}
        movt r0, #{hi(RING_BUF)}
        ubfx r1, r7, #0, #12
        ldrb r6, [r0, r1]
        adds r7, #1
        movw r4, #{lo(ST_STATE)}
        movt r4, #{hi(ST_STATE)}
        ldrb r5, [r4]
        cmp r5, #0
    """)
    a.b("s_wait_aa", EQ)
    a.asm("cmp r5, #1")
    a.b("s_len", EQ)
    # state 2: собираем в буфер
    a.asm(f"""
        movw r0, #{lo(ST_IDX)}
        movt r0, #{hi(ST_IDX)}
        ldrb r1, [r0]
        movw r2, #{lo(ST_BUF)}
        movt r2, #{hi(ST_BUF)}
        strb r6, [r2, r1]
        adds r1, #1
        strb r1, [r0]
        movw r0, #{lo(ST_LEN)}
        movt r0, #{hi(ST_LEN)}
        ldrb r0, [r0]
        cmp r1, r0
    """)
    a.b("ploop", LO)   # idx<len? продолжаем; когда idx==len — отправляем
    # idx==len -> отправить в USART3
    a.asm(f"""
        movw r0, #{lo(ST_BUF)}
        movt r0, #{hi(ST_BUF)}
        movw r1, #{lo(ST_LEN)}
        movt r1, #{hi(ST_LEN)}
        ldrb r1, [r1]
        bl #{USART3_TXBUF}
        movs r0, #0
        strb r0, [r4]
    """)
    a.b("ploop")
    # state 0: ждём 0xAA
    a.label("s_wait_aa")
    a.asm("cmp r6, #0xAA")
    a.b("ploop", NE)
    a.asm("movs r0, #1\nstrb r0, [r4]")
    a.b("ploop")
    # state 1: принять len
    a.label("s_len")
    a.asm(f"""
        movw r0, #{lo(ST_LEN)}
        movt r0, #{hi(ST_LEN)}
        strb r6, [r0]
        movw r0, #{lo(ST_IDX)}
        movt r0, #{hi(ST_IDX)}
        movs r1, #0
        strb r1, [r0]
        movs r0, #2
        strb r0, [r4]
    """)
    a.b("ploop")
    a.label("pret")
    a.asm(f"""
        movw r0, #{lo(MY_TAIL)}
        movt r0, #{hi(MY_TAIL)}
        str r7, [r0]
        pop {{r4, r5, r6, r7, pc}}
    """)
    poll=a.build()

    # ---- tick_router: init магии один раз + agent_poll каждый тик ----
    tick_at=poll_at+len(poll)
    r=Asm(ks,tick_at)
    r.asm(f"""
        push {{lr}}
        bl #{TICK_FUNC}
        movw r0, #{lo(MAGIC_ADDR)}
        movt r0, #{hi(MAGIC_ADDR)}
        ldr r1, [r0]
        movw r2, #{lo(MAGIC_VAL)}
        movt r2, #{hi(MAGIC_VAL)}
        cmp r1, r2
    """)
    r.b("inited", EQ)
    r.asm(f"""
        str r2, [r0]
        movw r0, #{lo(ST_STATE)}
        movt r0, #{hi(ST_STATE)}
        movs r1, #0
        strb r1, [r0]
        movw r0, #{lo(RING_HEAD)}
        movt r0, #{hi(RING_HEAD)}
        ldr r1, [r0]
        movw r0, #{lo(MY_TAIL)}
        movt r0, #{hi(MY_TAIL)}
        str r1, [r0]
    """)
    r.label("inited")
    r.asm(f"bl #{poll_at | 1}\npop {{pc}}")
    tick=r.build()

    # ---- снифер on_frame ----
    hook_at=tick_at+len(tick)
    h=Asm(ks,hook_at)
    h.asm(f"""
        push {{r4, r5, r6, lr}}
        mov r4, r0
        ldrb r5, [r4, #6]
        ldrb r6, [r4, #7]
        orr.w r5, r6, r5, lsl #8
        adds r5, #9
        uxtb r5, r5
        movs r0, #0xF1
        bl #{TX_BYTE}
        movs r0, #0xF2
        bl #{TX_BYTE}
        mov r0, r5
        bl #{TX_BYTE}
        movs r6, #0
    """)
    h.label("sloop")
    h.asm("cmp r6, r5")
    h.b("sdone", HS)
    h.asm(f"ldrb r0, [r4, r6]\nbl #{TX_BYTE}\nadds r6, #1")
    h.b("sloop")
    h.label("sdone")
    h.asm(f"mov r0, r4\nbl #{ON_FRAME}\npop {{r4, r5, r6, pc}}")
    hook=h.build()

    code=poll+tick+hook
    img+=code
    while len(img)%4: img.append(0xFF)
    print(f"poll 0x{poll_at:08X} ({len(poll)}Б), tick 0x{tick_at:08X}, hook 0x{hook_at:08X}")

    img[UI_HOOK-BASE:UI_HOOK-BASE+4]=thumb_bl(UI_HOOK,tick_at)
    img[FRAME_CALL-BASE:FRAME_CALL-BASE+4]=thumb_bl(FRAME_CALL,hook_at)
    open(OUT,"wb").write(bytes(img))
    print(f"готово: {OUT}, {len(img)} Б")


if __name__=="__main__":
    main()