from machine import Pin, UART
import struct
import time
import sys
import select


# ============================================================
# ESP32 + IMU + USB binary bridge
#
# Wiring:
#   IMU TX yellow -> ESP32 GPIO16 / RX2
#   IMU RX green  -> ESP32 GPIO17 / TX2
#   IMU GND       -> ESP32 GND
#
# UART2: IMU
# USB stdout/stdin: PC
#
# Important:
#   - Output uses sys.stdout, because this is the stable path.
#   - Command input uses sys.stdin, ASCII only: o / c / s.
#   - First test with USE_REAL_DOUT = False.
# ============================================================


# ---------------- PIN CONFIG ----------------

IMU_RX_PIN = 16
IMU_TX_PIN = 17

DIN_PINS = (25, 26, 27, 32)
DOUT_PINS = (18, 19)

digital_inputs = [Pin(pin, Pin.IN, Pin.PULL_UP) for pin in DIN_PINS]
digital_outputs = [Pin(pin, Pin.OUT, value=0) for pin in DOUT_PINS]


# ---------------- SERIAL CONFIG ----------------

# PC side uses MicroPython stdout/stdin.
# PC reader should use:
#   python3 pc_reader.py /dev/ttyUSB0 --baud 115200
PC_BAUD_NOTE = 115200

# IMU side confirmed working.
IMU_BAUD = 921600

SEND_PERIOD_MS = 20      # 50 Hz
STARTUP_DELAY_MS = 3000  # quiet window after reset


# ---------------- IMU FRAME CONFIG ----------------

FRAME_LEN = 11
HEADER = 0x55

TYPE_ACC = 0x51
TYPE_GYRO = 0x52
TYPE_ANGLE = 0x53

# raw_values:
# acc x y z, gyro x y z, angle x y z, temperature
raw_values = [0] * 10


# ---------------- DOUT TEST MODE ----------------

# False:
#   Only update software DOUT state in outgoing packet.
#   GPIO18/GPIO19 will NOT really change.
#
# True:
#   Actually drive GPIO18/GPIO19.
#
# Start with False. If o/c/s changes DOUT in pc_reader output,
# then set it to True and test real hardware.
USE_REAL_DOUT = True

# bit0 = DOUT0 / GPIO18
# bit1 = DOUT1 / GPIO19
dout_shadow = 0


# ---------------- BUFFERS ----------------

rx_buffer = bytearray(2048)
write_idx = 0
read_idx = 0


# ---------------- USB STDOUT / STDIN ----------------

try:
    usb_out = sys.stdout.buffer
except AttributeError:
    usb_out = sys.stdout

poller = select.poll()
poller.register(sys.stdin, select.POLLIN)


# ============================================================
# DOUT / gripper
# ============================================================

def clear_douts():
    global dout_shadow

    dout_shadow = 0

    for pin in digital_outputs:
        pin.value(0)


def set_gripper(mode):
    global dout_shadow

    # mode:
    #   0 = STOP
    #   1 = OPEN
    #   2 = CLOSE
    #
    # DOUT[0] = GPIO18
    # DOUT[1] = GPIO19

    if mode == 1:
        # OPEN -> DOUT=[0, 1]
        dout_shadow = 0b10

    elif mode == 2:
        # CLOSE -> DOUT=[1, 0]
        dout_shadow = 0b01

    else:
        # STOP -> DOUT=[0, 0]
        dout_shadow = 0b00

    # Shadow test mode: do not touch real GPIO pins.
    if not USE_REAL_DOUT:
        return

    # Real output mode.
    digital_outputs[0].value(0)
    digital_outputs[1].value(0)
    time.sleep_ms(20)

    if mode == 1:
        # OPEN
        digital_outputs[1].value(1)

    elif mode == 2:
        # CLOSE
        digital_outputs[0].value(1)

    # mode 0 keeps both LOW


# ============================================================
# IMU UART
# ============================================================

def open_imu_uart():
    uart = UART(
        2,
        baudrate=IMU_BAUD,
        bits=8,
        parity=None,
        stop=1,
        rx=IMU_RX_PIN,
        tx=IMU_TX_PIN,
        timeout=0,
        rxbuf=4096,
    )

    time.sleep_ms(300)

    # Flush startup junk.
    while uart.any():
        uart.read()

    return uart


# ============================================================
# IMU parser
# ============================================================

def update_raw_values(frame):
    frame_type = frame[1]
    x, y, z, extra = struct.unpack_from("<hhhh", frame, 2)

    if frame_type == TYPE_ACC:
        raw_values[0:3] = (x, y, z)
        raw_values[9] = extra

    elif frame_type == TYPE_GYRO:
        raw_values[3:6] = (x, y, z)

    elif frame_type == TYPE_ANGLE:
        raw_values[6:9] = (x, y, z)
        return True

    return False


def poll_imu(uart):
    global write_idx, read_idx

    count = uart.any()
    if count:
        data = uart.read(count)
        if data:
            for b in data:
                rx_buffer[write_idx] = b
                write_idx = (write_idx + 1) % len(rx_buffer)

                # Buffer full: drop oldest byte.
                if write_idx == read_idx:
                    read_idx = (read_idx + 1) % len(rx_buffer)

    updated = False

    while True:
        available = (write_idx - read_idx) % len(rx_buffer)
        if available < FRAME_LEN:
            break

        # Resync to 0x55.
        if rx_buffer[read_idx] != HEADER:
            read_idx = (read_idx + 1) % len(rx_buffer)
            continue

        frame = bytearray(FRAME_LEN)
        for i in range(FRAME_LEN):
            frame[i] = rx_buffer[(read_idx + i) % len(rx_buffer)]

        # Checksum.
        if (sum(frame[:10]) & 0xFF) != frame[10]:
            read_idx = (read_idx + 1) % len(rx_buffer)
            continue

        read_idx = (read_idx + FRAME_LEN) % len(rx_buffer)

        if update_raw_values(frame):
            updated = True

    return updated


# ============================================================
# Packet to PC
# ============================================================

def make_packet():
    din_mask = 0
    for index, pin in enumerate(digital_inputs):
        din_mask |= pin.value() << index

    if USE_REAL_DOUT:
        dout_mask = 0
        for index, pin in enumerate(digital_outputs):
            dout_mask |= pin.value() << index
    else:
        dout_mask = dout_shadow

    # Packet layout:
    # A5 5A | version | DIN | DOUT | ticks_ms | 10 signed int16 | checksum
    packet = struct.pack(
        "<2sBBBI10h",
        b"\xA5\x5A",
        2,
        din_mask,
        dout_mask,
        time.ticks_ms() & 0xFFFFFFFF,
        *raw_values,
    )

    return packet + bytes((sum(packet) & 0xFF,))


def usb_write(data):
    usb_out.write(data)


# ============================================================
# USB command parser
# ============================================================

def poll_usb_commands():
    # ASCII only.
    # pc_reader.py should send:
    #   b"o" = OPEN
    #   b"c" = CLOSE
    #   b"s" = STOP
    #
    # Do NOT use C3 3C binary command here.

    for _ in range(8):
        if not poller.poll(0):
            break

        ch = sys.stdin.read(1)

        if not ch:
            break

        if ch == "o" or ch == "O":
            set_gripper(1)

        elif ch == "c" or ch == "C":
            set_gripper(2)

        elif ch == "s" or ch == "S":
            set_gripper(0)

        # Ignore Enter/newline and other characters.


# ============================================================
# Main
# ============================================================

clear_douts()

# Quiet window after reset.
# During this period there is no binary output.
time.sleep_ms(STARTUP_DELAY_MS)

imu_uart = open_imu_uart()

last_send_ms = time.ticks_ms()

while True:
    poll_usb_commands()
    poll_imu(imu_uart)

    now = time.ticks_ms()

    # Fixed 50 Hz output.
    if time.ticks_diff(now, last_send_ms) >= SEND_PERIOD_MS:
        usb_write(make_packet())
        last_send_ms = now

    time.sleep_ms(0)