From cf06b97524ccede8fd688406bd66e9ecc6b5ff3b Mon Sep 17 00:00:00 2001
From: vanminhpe <vanminh@xnergytech.com>
Date: Mon, 27 Jul 2026 11:25:54 +0800
Subject: [PATCH] fix: wait for TPU allowed state before triggering flash write

---
 test_modbus_limiter.py | 43 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/test_modbus_limiter.py b/test_modbus_limiter.py
index c67f06c..ed5d34f 100644
--- a/test_modbus_limiter.py
+++ b/test_modbus_limiter.py
@@ -62,6 +62,17 @@ PASSWORD_HI      = 0xCEE5
 # TPU trigger: holding coil bit 49 (fenghuang-dsp modbus_mapping.h)
 TPU_BIT_WRITE_FLASH = 49   # MODBUS_ALL_CFG_CMD_BIT → main_app.write_flash
 
+# TPU state register: input reg 61 (MODBUS_MAIN_STATE), read via FC4.
+# Firmware only processes write_flash in these states (app_modbus.c update_modbus_input).
+ADDR_TPU_STATE = 61
+TPU_TRIGGER_ALLOWED_STATES = {
+    2,    # TPU_READY
+    10,   # TPU_INITIALIZATION
+    15,   # TPU_DEBUG_CONFIG
+    16,   # TPU_TEST_MODE
+    0xEE, # TPU_ERROR
+}
+
 # Auto-detection register (holding reg 163, same address for both RPU and TPU)
 ADDR_HARDWARE_REV_HI = 163  # MODBUS_HARDWARE_REVISION_HI
 HW_REV_TPU = 0x0021
@@ -310,7 +321,7 @@ def read_reg(instr: minimalmodbus.Instrument, addr: int, data_type: str):
 
 def write_reg(instr: minimalmodbus.Instrument, addr: int, value, data_type: str) -> None:
     if _is_float(data_type):
-        instr.write_float(addr, float(value), functioncode=16, number_of_registers=2,
+        instr.write_float(addr, float(value), number_of_registers=2,
                           byteorder=minimalmodbus.BYTEORDER_BIG)
         return
     signed = (data_type == 'int16')
@@ -352,8 +363,31 @@ def trigger_rewrite_rpu(instr: minimalmodbus.Instrument) -> None:
     _wait_for_device(instr, ADDR_PASSWORD_LO)
 
 
+def _wait_for_tpu_trigger_state(instr: minimalmodbus.Instrument) -> bool:
+    """Poll input reg 61 (MODBUS_MAIN_STATE) until state allows write_flash trigger.
+
+    Returns True if state reached, False if timed out.
+    Required because firmware only processes write_flash in specific states.
+    """
+    deadline = time.monotonic() + REWRITE_WAIT
+    while time.monotonic() < deadline:
+        try:
+            state = instr.read_register(ADDR_TPU_STATE, number_of_decimals=0,
+                                        functioncode=4, signed=False)
+            if state in TPU_TRIGGER_ALLOWED_STATES:
+                print(f"  (trigger: TPU state=0x{state:02X} — ready to accept trigger)")
+                return True
+            print(f"  (trigger: TPU state=0x{state:02X} — waiting for allowed state...)")
+        except minimalmodbus.ModbusException:
+            pass
+        time.sleep(0.5)
+    print("  (trigger: WARNING — TPU state never reached allowed state; trigger may be ignored)")
+    return False
+
+
 def trigger_rewrite_tpu(instr: minimalmodbus.Instrument) -> None:
     """Write coil bit 49 (MODBUS_ALL_CFG_CMD_BIT) → modbus_to_config() + write_all_flash() + soft_reset()."""
+    _wait_for_tpu_trigger_state(instr)
     try:
         instr.write_bit(TPU_BIT_WRITE_FLASH, 1, functioncode=5)
         print("  (trigger: coil bit write ack received)")
@@ -392,9 +426,11 @@ def run_tests(port: str) -> tuple[str, list[dict]]:
     if device_type == 'TPU':
         TESTS = TPU_TESTS
         trigger_fn = trigger_rewrite_tpu
+        time.sleep(5)
     else:
         TESTS = RPU_TESTS
         trigger_fn = trigger_rewrite_rpu
+        time.sleep(5)
     print(f"  Device type: {device_type} — {len(TESTS)} registers in test list")
 
     # -------------------------------------------------------------------------
@@ -427,7 +463,8 @@ def run_tests(port: str) -> tuple[str, list[dict]]:
 
     print("\n[Phase 1] Triggering rewrite...")
     trigger_fn(instr)
-
+    time.sleep(5)
+    
     print("[Phase 1] Reading back...")
     below_readbacks: dict[int, any] = {}
     for (_, addr, data_type, *_) in TESTS:
@@ -452,6 +489,7 @@ def run_tests(port: str) -> tuple[str, list[dict]]:
 
     print("\n[Phase 2] Triggering rewrite...")
     trigger_fn(instr)
+    time.sleep(5)
 
     print("[Phase 2] Reading back...")
     above_readbacks: dict[int, any] = {}
@@ -474,6 +512,7 @@ def run_tests(port: str) -> tuple[str, list[dict]]:
             except minimalmodbus.ModbusException as e:
                 print(f"  addr {addr:3d}: restore write ERROR — {e}")
     trigger_fn(instr)
+    time.sleep(5)
     print("[Restore] Done.")
 
     # -------------------------------------------------------------------------