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
#
# Physical switch:
#   GPIO27 S -> toggle switch one side
#   GPIO27 G -> toggle switch other side
#
# GPIO27 switch behavior:
#   OFF -> ON  : CLOSE pulse
#   ON  -> OFF : OPEN pulse
# ============================================================
# 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]

# Low-level trigger relay:
# GPIO HIGH = relay OFF
# GPIO LOW  = relay ON
digital_outputs = [Pin(pin, Pin.OUT, value=1) for pin in DOUT_PINS]


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

PC_BAUD_NOTE = 115200
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


# ---------------- PHYSICAL SWITCH CONFIG ----------------

# DIN_PINS = (25, 26, 27, 32)
# GPIO27 is index 2.
SWITCH_INDEX = 2

# How long each switch movement triggers the relay.
# If the cutter movement is too short, increase to 500 / 800.
SWITCH_PULSE_MS = 2000

# Debounce time for mechanical toggle switch.
SWITCH_DEBOUNCE_MS = 80

# Initial switch state.
# PULL_UP:
#   switch open  = 1
#   switch closed = 0
switch_last_value = digital_inputs[SWITCH_INDEX].value()
switch_last_edge_ms = time.ticks_ms()

switch_pulse_active = False
switch_pulse_start_ms = time.ticks_ms()


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

    # Logical DOUT state reported to PC.
    # 0 means no gripper command active.
    dout_shadow = 0

    # Low-level trigger relay:
    # HIGH = relay OFF
    # LOW  = relay ON
    for pin in digital_outputs:
        pin.value(1)


def set_gripper(mode):
    global dout_shadow

    # mode:
    #   0 = STOP
    #   1 = OPEN
    #   2 = CLOSE
    #
    # DOUT[0] = GPIO18 -> relay IN1
    # DOUT[1] = GPIO19 -> relay IN2
    #
    # Low-level trigger relay:
    #   GPIO HIGH = relay OFF
    #   GPIO LOW  = relay ON
    #
    # Tested hardware behavior:
    #   IN1 = LOW, IN2 = HIGH -> CLOSE
    #   IN1 = HIGH, IN2 = LOW -> OPEN

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

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

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

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

    # First turn both relays OFF.
    # This prevents IN1 and IN2 being active at the same time.
    digital_outputs[0].value(1)  # GPIO18 / IN1 OFF
    digital_outputs[1].value(1)  # GPIO19 / IN2 OFF
    time.sleep_ms(20)

    if mode == 1:
        # OPEN -> trigger IN2 / GPIO19
        digital_outputs[1].value(0)

    elif mode == 2:
        # CLOSE -> trigger IN1 / GPIO18
        digital_outputs[0].value(0)

    # mode 0 keeps both HIGH/OFF


# ============================================================
# Physical switch control
# ============================================================

def start_switch_pulse(mode):
    global switch_pulse_active, switch_pulse_start_ms

    set_gripper(mode)
    switch_pulse_active = True
    switch_pulse_start_ms = time.ticks_ms()


def poll_physical_switch():
    global switch_last_value, switch_last_edge_ms

    now = time.ticks_ms()
    value = digital_inputs[SWITCH_INDEX].value()  # GPIO27

    # Detect state change with debounce.
    if value != switch_last_value:
        if time.ticks_diff(now, switch_last_edge_ms) > SWITCH_DEBOUNCE_MS:
            old_value = switch_last_value
            switch_last_value = value
            switch_last_edge_ms = now

            # OFF -> ON:
            # GPIO27 goes 1 -> 0
            # switch connects S to G
            if old_value == 1 and value == 0:
                start_switch_pulse(2)   # CLOSE

            # ON -> OFF:
            # GPIO27 goes 0 -> 1
            elif old_value == 0 and value == 1:
                start_switch_pulse(1)   # OPEN


def update_switch_pulse_stop():
    global switch_pulse_active

    if not switch_pulse_active:
        return

    now = time.ticks_ms()

    if time.ticks_diff(now, switch_pulse_start_ms) >= SWITCH_PULSE_MS:
        set_gripper(0)   # STOP
        switch_pulse_active = False


# ============================================================
# 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, but do NOT wait forever.
    # IMU is a continuous stream; an unlimited while uart.any() may get stuck.
    t0 = time.ticks_ms()
    while time.ticks_diff(time.ticks_ms(), t0) < 100:
        if uart.any():
            uart.read()
        else:
            time.sleep_ms(1)

    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

    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)
    try:
        usb_out.flush()
    except AttributeError:
        pass


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

def poll_usb_commands():
    global switch_pulse_active

    # ASCII only.
    # pc_reader.py should send:
    #   b"o" = OPEN
    #   b"c" = CLOSE
    #   b"s" = STOP

    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":
            switch_pulse_active = False
            set_gripper(1)

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

        elif ch == "s" or ch == "S":
            switch_pulse_active = False
            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_physical_switch()
    update_switch_pulse_stop()
    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)