import time

from common.config_loader import load_fenghuang_registers
from drivers.modbus_device import ModbusDevice, ModbusDeviceError

DEFAULT_DATASETS = ["Pri_normal", "Debug", "event_log"]

# TPU_STATE values that allow modbus_to_config()/write_all_flash() to run
# (fenghuang-dsp/app/app_modbus/app_modbus.c update_modbus_input()).
# TPU_RUN_INV (charging) is NOT in this set - config writes are silently
# ignored while charging is active.
CONFIG_WRITABLE_STATES = {2, 10, 15, 16, 238}  # READY, INITIALIZATION, DEBUG_CONFIG, TEST_MODE, ERROR


class Tpu(ModbusDevice):
    """Covers every register in Pri_normal + Debug + event_log by label via
    read()/write() (inherited from ModbusDevice). More datasets can be added
    per-test via extra_datasets. Only registers that need real logic beyond a
    1:1 read/write get a dedicated method below.
    """

    REWRITE_WAIT_S = 5

    def __init__(self, port, extra_datasets=None, slave_id=0x10):
        registers = load_fenghuang_registers(DEFAULT_DATASETS + list(extra_datasets or []))
        super().__init__(port, "TPU", registers, slave_id=slave_id)

    def wait_until_config_writable(self, timeout_s=10, poll_interval_s=0.5):
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            if self.read("Main_state") in CONFIG_WRITABLE_STATES:
                return
            time.sleep(poll_interval_s)
        raise ModbusDeviceError(f"TPU did not reach a config-writable state within {timeout_s}s")

    def read_error_code(self):
        return (self.read("Error_code_HI") << 16) | self.read("Error_code_LO")

    def read_shadow_error_code(self):
        return (self.read("Shadow_Error_code_HI") << 16) | self.read("Shadow_Error_code_LO")

    def rewrite_config(self):
        self.write_coil(self._addr("rewrite_cfg_flash"), True, verify=False)
        time.sleep(self.REWRITE_WAIT_S)