import struct
import time
import serial


SYNC = b"\xA5\x5A"
PACKET_FORMAT = "<2sBBBI10hB"
PACKET_SIZE = struct.calcsize(PACKET_FORMAT)


def decode_packet(packet):
    fields = struct.unpack(PACKET_FORMAT, packet)
    _, version, din_mask, dout_mask, time_ms, *rest = fields
    raw = rest[:-1]

    acc_scale = 16.0 / 32768.0
    gyro_scale = 2000.0 / 32768.0
    angle_scale = 180.0 / 32768.0

    return {
        "version": version,
        "t_ms": time_ms,
        "din": [(din_mask >> bit) & 1 for bit in range(4)],
        "dout": [(dout_mask >> bit) & 1 for bit in range(2)],
        "acc_g": [round(value * acc_scale, 5) for value in raw[0:3]],
        "gyro_dps": [round(value * gyro_scale, 4) for value in raw[3:6]],
        "angle_deg": [round(value * angle_scale, 4) for value in raw[6:9]],
        "temp_c": round(raw[9] / 100.0, 2),
    }


class ESP32Bridge:
    def __init__(
        self,
        port,
        baud=115200,
        timeout=0.05,
        startup_wait=1.0,
        safe_stop_on_close=True,
        reset_input_on_open=True,
        reset_output_on_open=True,
    ):
        self.port = port
        self.baud = baud
        self.timeout = timeout
        self.startup_wait = startup_wait
        self.safe_stop_on_close = safe_stop_on_close
        self.reset_input_on_open = reset_input_on_open
        self.reset_output_on_open = reset_output_on_open

        self.ser = None
        self.buffer = bytearray()
        self.last_t_ms = None

    def is_open(self):
        return self.ser is not None and self.ser.is_open

    def _require_open(self):
        if not self.is_open():
            raise RuntimeError("Serial port is not open.")

    def open(self):
        if self.is_open():
            return self

        self.ser = serial.Serial()
        self.ser.port = self.port
        self.ser.baudrate = self.baud
        self.ser.timeout = self.timeout
        self.ser.write_timeout = self.timeout

        # Avoid ESP32 auto reset / boot mode weirdness.
        self.ser.dtr = False
        self.ser.rts = False
        self.ser.open()

        self.ser.dtr = False
        self.ser.rts = False

        if self.reset_input_on_open:
            self.ser.reset_input_buffer()

        if self.reset_output_on_open:
            self.ser.reset_output_buffer()

        self.buffer.clear()
        self.last_t_ms = None

        if self.startup_wait > 0:
            time.sleep(self.startup_wait)

        return self

    def close(self):
        if self.ser is None:
            return

        if self.safe_stop_on_close and self.ser.is_open:
            try:
                self.stop_gripper()
                time.sleep(0.1)
            except Exception:
                pass

        try:
            if self.ser.is_open:
                self.ser.close()
        finally:
            self.ser = None

    def __enter__(self):
        return self.open()

    def __exit__(self, exc_type, exc, tb):
        self.close()

    def send_gripper_command(self, mode):
        self._require_open()

        if mode == 1:
            self.ser.write(b"o")
        elif mode == 2:
            self.ser.write(b"c")
        else:
            self.ser.write(b"s")

        self.ser.flush()

    def open_gripper(self):
        self.send_gripper_command(1)

    def close_gripper(self):
        self.send_gripper_command(2)

    def stop_gripper(self):
        self.send_gripper_command(0)

    def read_packet(self):
        """
        Return one decoded packet dict.
        This blocks until a valid packet is found.
        """
        self._require_open()

        while True:
            chunk = self.ser.read(self.ser.in_waiting or 1)
            if chunk:
                self.buffer.extend(chunk)

            packet = self._try_pop_packet()
            if packet is not None:
                return decode_packet(packet)

    def _try_pop_packet(self):
        while len(self.buffer) >= PACKET_SIZE:
            sync_index = self.buffer.find(SYNC)

            if sync_index < 0:
                # Keep last byte in case it is the first half of SYNC.
                del self.buffer[:-1]
                return None

            if sync_index:
                del self.buffer[:sync_index]

            if len(self.buffer) < PACKET_SIZE:
                return None

            packet = bytes(self.buffer[:PACKET_SIZE])

            if (sum(packet[:-1]) & 0xFF) != packet[-1]:
                del self.buffer[0]
                continue

            del self.buffer[:PACKET_SIZE]
            return packet

        return None

    def read_packet_with_frequency(self):
        data = self.read_packet()

        if self.last_t_ms is None:
            freq_hz = 0.0
        else:
            delta = (data["t_ms"] - self.last_t_ms) & 0xFFFFFFFF
            freq_hz = 1000.0 / delta if delta else 0.0

        self.last_t_ms = data["t_ms"]
        data["freq_hz"] = freq_hz
        return data

    def iter_packets(self, seconds=None, max_samples=None, with_frequency=True):
        start = time.monotonic()
        count = 0

        while True:
            if seconds is not None and time.monotonic() - start >= seconds:
                break

            if max_samples is not None and count >= max_samples:
                break

            if with_frequency:
                yield self.read_packet_with_frequency()
            else:
                yield self.read_packet()

            count += 1