import time

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

DEFAULT_DATASETS = ["Sec_normal", "sec_flash"]


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 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)