import time

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

DEFAULT_DATASETS = ["Sec_normal", "sec_flash"]

# SEC_STATE values that allow modbus_to_config()/write_all_flash() to run
# (garuda-dsp/app/app_modbus/app_modbus.c update_modbus_input()).
# SEC_RUN (charging) is NOT in this set - config writes are silently ignored
# while charging is active.
CONFIG_WRITABLE_STATES = {1, 9, 238, 240, 241}  # READY, PRE_READY, ERROR, DEBUG_READY, DEBUG_FREQ_CAL


class Rpu(ModbusDevice):
    """Covers every register in Sec_normal + sec_flash by label via read()/write()
    (inherited from ModbusDevice). Extra 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.
    """

    Q128_DIVISOR = 128.0
    REWRITE_PASSWORD_LO = 0xFEED
    REWRITE_PASSWORD_HI = 0xCEE5
    REWRITE_WAIT_S = 5

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

    def set_charging(self, enable):
        self.write_coil(self._addr("Enable_charger"), enable)

    def set_runtime_voltage(self, volts):
        self.write_holding_register(self._addr("Runtime_voltage_setting_Q128"), round(volts * self.Q128_DIVISOR))

    def set_runtime_current(self, amps):
        self.write_holding_register(self._addr("Runtime_current_setting_Q128"), round(amps * self.Q128_DIVISOR))

    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 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"RPU did not reach a config-writable state within {timeout_s}s")

    def rewrite_config(self):
        self.write_holding_register(self._addr("PASSWORD_LO"), self.REWRITE_PASSWORD_LO)
        self.write_holding_register(self._addr("PASSWORD_HI"), self.REWRITE_PASSWORD_HI)
        time.sleep(self.REWRITE_WAIT_S)