"""Fast MCU client using native LowState decrypt/parse and UDP drain.

The public shape mirrors go1_pro_sdk.connection.MCUClient for deployment code,
while keeping the receive hot path out of pure Python.
"""
import time

from fast_lowcmd import FastLowCmdBuilder, default_state_path
from go1_fast_lowcmd import FastNativeMCUClient
from go1_pro_sdk import MCU_IP, MCU_PORT
from go1_pro_sdk.utils.constants import RCVBUF_SIZE


class FastIMU:
    __slots__ = ("quaternion", "gyroscope", "accelerometer", "rpy", "temperature")

    def __init__(self, fields):
        self.quaternion = tuple(fields["imu_quaternion"])
        self.gyroscope = tuple(fields["imu_gyroscope"])
        self.accelerometer = (0.0, 0.0, 0.0)
        self.rpy = tuple(fields["imu_rpy"])
        self.temperature = 0


class FastMotorState:
    __slots__ = ("mode", "q", "dq", "ddq", "tauEst", "q_raw", "dq_raw", "ddq_raw", "temperature", "reserve")

    def __init__(self, fields, i):
        self.mode = int(fields["motor_mode"][i])
        self.q = float(fields["motor_q"][i])
        self.dq = float(fields["motor_dq"][i])
        self.ddq = float(fields["motor_ddq"][i])
        self.tauEst = float(fields["motor_tau"][i])
        self.q_raw = float(fields["motor_q_raw"][i])
        self.dq_raw = float(fields["motor_dq_raw"][i])
        self.ddq_raw = float(fields["motor_ddq_raw"][i])
        self.temperature = int(fields["motor_temperature"][i])
        self.reserve = [int(fields["motor_reserve0"][i]), int(fields["motor_reserve1"][i])]


class FastBMS:
    __slots__ = (
        "version_h", "version_l", "bms_status", "SOC", "current", "cycle",
        "BQ_NTC", "MCU_NTC", "cell_vol",
    )

    def __init__(self, fields):
        self.version_h = int(fields["bms_version_h"])
        self.version_l = int(fields["bms_version_l"])
        self.bms_status = int(fields["bms_status"])
        self.SOC = int(fields["bms_soc"])
        self.current = int(fields["bms_current"])
        self.cycle = int(fields["bms_cycle"])
        self.BQ_NTC = [int(fields["bms_bq_ntc0"]), int(fields["bms_bq_ntc1"])]
        self.MCU_NTC = [int(fields["bms_mcu_ntc0"]), int(fields["bms_mcu_ntc1"])]
        self.cell_vol = [int(x) for x in fields["bms_cell_vol"]]

    @property
    def voltage_mv(self):
        return sum(self.cell_vol)

    @property
    def voltage_v(self):
        return self.voltage_mv / 1000.0

    @property
    def current_a(self):
        return self.current / 1000.0


class FastRemoteState:
    __slots__ = ("head", "btn", "lx", "ly", "rx", "ry", "L2", "pressed")

    def __init__(self, fields):
        axes = fields["remote_axes"]
        self.head = b"\x55\xaa"
        self.btn = int(fields["remote_btn"])
        self.lx = float(axes[0])
        self.ly = float(axes[1])
        self.rx = float(axes[2])
        self.ry = float(axes[3])
        self.L2 = float(axes[4])
        self.pressed = list(fields["remote_pressed"])

    def is_pressed(self, name):
        return name in self.pressed

    def any_button(self):
        return self.btn != 0


class FastLowState:
    __slots__ = (
        "head", "levelFlag", "frameReserve", "SN", "version", "bandWidth",
        "imu", "motorState", "footForce", "footForceEst", "bms", "tick",
        "wirelessRemote", "reserve", "crc", "remote",
    )

    def __init__(self, fields):
        self.head = int(fields["head"])
        self.levelFlag = int(fields["levelFlag"])
        self.frameReserve = int(fields["frameReserve"])
        self.SN = b"\x00" * 8
        self.version = b"\x00" * 8
        self.bandWidth = int(fields["bandWidth"])
        self.imu = FastIMU(fields)
        self.motorState = [FastMotorState(fields, i) for i in range(20)]
        self.footForce = (0, 0, 0, 0)
        self.footForceEst = (0, 0, 0, 0)
        self.bms = FastBMS(fields)
        self.tick = 0
        self.wirelessRemote = b"\x00" * 40
        self.reserve = b"\x00" * 4
        self.crc = b"\x00" * 4
        self.remote = FastRemoteState(fields)


class FastMCUClient:
    def __init__(
            self,
            state_path=None,
            mcu_ip=MCU_IP,
            mcu_port=MCU_PORT,
            local_port=0,
            endian="little"):
        if endian != "little":
            raise ValueError("FastMCUClient currently supports only little-endian Blowfish state")
        self.mcu_ip = mcu_ip
        self.mcu_port = mcu_port
        resolved_state_path = state_path or default_state_path()
        self.builder = FastLowCmdBuilder(resolved_state_path)
        self._native = FastNativeMCUClient(
            resolved_state_path,
            mcu_ip,
            int(mcu_port),
            int(local_port),
            int(RCVBUF_SIZE),
            endian,
        )
        self.local_port = self._native.local_port()
        self._last_state = None
        self.backend = "cpp_lowcmd_cpp_lowstate_cpp_udp"

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def close(self):
        native = getattr(self, "_native", None)
        if native is not None:
            native.close()
            self._native = None

    def send_raw(self, raw_cipher):
        return self._native.send_raw(raw_cipher)

    def recv_latest(self):
        fields = self._native.recv_latest_fields()
        if fields is None:
            return None
        self._last_state = FastLowState(fields)
        return self._last_state

    def recv_state(self, timeout=1.0):
        deadline = time.time() + timeout
        while time.time() < deadline:
            state = self.recv_latest()
            if state is not None:
                return state
            time.sleep(0.001)
        return None

    @property
    def last_state(self):
        return self._last_state

    def wake_mcu(self, n_frames=50, dt=0.01):
        damping = self.builder.build_encrypted_damping()
        recv_count = 0
        for _ in range(n_frames):
            self.send_raw(damping)
            time.sleep(dt)
            if self.recv_latest() is not None:
                recv_count += 1
        return recv_count

    def safe_stop(self, n_frames=50, dt=0.002):
        damping = self.builder.build_encrypted_damping()
        for _ in range(n_frames):
            try:
                self.send_raw(damping)
            except Exception:
                pass
            time.sleep(dt)