from machine import UART
import time
import struct
import gc

# ================= 設定區 =================
FC_BAUDRATE = 115200 
XB_BAUDRATE = 115200
DEST_64 = b'\x00\x13\xA2\x00\x42\x5B\x9D\xC8' 
XBEE_MAX_PAYLOAD = 100
MY_SYSID = 15  # 第二台就改成 10，第三台 15
POLL_MAGIC = b'POLL'

tx_buf = bytearray()
MAX_BUF_SIZE = 6144 # 保留安全餘裕的緩衝區大小
# =========================================

uart_fc = UART(1, baudrate=FC_BAUDRATE, tx=32, rx=33, rxbuf=4096)
uart_xb = UART(2, baudrate=XB_BAUDRATE, tx=25, rx=26, rxbuf=4096)

rx_buf = bytearray()

def get_checksum(data):
    return 0xFF - (sum(data) & 0xFF)

def send_to_xbee_chunked(payload):
    total_len = len(payload)
    sent_len = 0
    
    while sent_len < total_len:
        end_len = min(sent_len + XBEE_MAX_PAYLOAD, total_len)
        chunk = payload[sent_len : end_len]
        sent_len = end_len
        
        frame_content = bytearray()
        frame_content.append(0x10)       
        frame_content.append(0x00) # 0x00 代表不要求 XBee ACK，極大加快發送速度       
        frame_content.extend(DEST_64)    
        frame_content.extend(b'\xFF\xFE')
        frame_content.append(0x00)       
        frame_content.append(0x00)       
        frame_content.extend(chunk)      
        
        length = len(frame_content)
        checksum = get_checksum(frame_content)
        
        packet = bytearray()
        packet.append(0x7E)
        packet.append((length >> 8) & 0xFF)
        packet.append(length & 0xFF)
        packet.extend(frame_content)
        packet.append(checksum)
        
        uart_xb.write(packet)
        # 僅需極短暫延遲，讓硬體發送緩衝區消化即可
        time.sleep_ms(1)

def process_xbee_buffer():
    global rx_buf, tx_buf
    while True:
        start_pos = rx_buf.find(b'\x7E')
        if start_pos == -1:
            rx_buf = bytearray() # 避免用 clear()
            return 
            
        if start_pos > 0:
            rx_buf = rx_buf[start_pos:]
            
        if len(rx_buf) < 3: return
            
        pkt_len = (rx_buf[1] << 8) | rx_buf[2]
        total_len = pkt_len + 4 
        
        if pkt_len > 300:
            rx_buf = rx_buf[1:]
            continue
        
        if len(rx_buf) < total_len: return 
            
        packet = rx_buf[3 : 3+pkt_len]
        checksum_recv = rx_buf[3+pkt_len]
        
        if get_checksum(packet) == checksum_recv:
            if packet[0] == 0x90:
                real_data = packet[12:]
                if len(real_data) == 5 and real_data.startswith(POLL_MAGIC):
                    if real_data[4] == MY_SYSID:
                        flush_tx_buffer()
                else:
                    uart_fc.write(real_data)
            
            rx_buf = rx_buf[total_len:]
        else:
            rx_buf = rx_buf[1:]

def flush_tx_buffer():
    global tx_buf
    if len(tx_buf) == 0:
        return
        
    # 單次最大發送量設為 1200 bytes，避免在空中飛行超過時槽導致碰撞
    send_limit = min(len(tx_buf), 1200) 
    
    cut_idx = send_limit
    if send_limit < len(tx_buf):
        last_fd_pos = tx_buf[:send_limit].rfind(b'\xFD')
        if last_fd_pos > 0:
            cut_idx = last_fd_pos 
            
    data_to_send = tx_buf[:cut_idx]
    tx_buf = tx_buf[cut_idx:] 
    
    send_to_xbee_chunked(data_to_send)

# 初始清理
gc.collect()

while True:
    try:
        # FC -> 緩衝區
        if uart_fc.any():
            data = uart_fc.read(250)
            if data:
                tx_buf.extend(data)
                
                # 聰明的溫和防爆機制 (砍掉舊的 2/3，保留最新 1/3)
                if len(tx_buf) > MAX_BUF_SIZE:
                    safe_cut = tx_buf.find(b'\xFD', (MAX_BUF_SIZE // 3) * 2)
                    if safe_cut != -1:
                        tx_buf = tx_buf[safe_cut:] 
                    else:
                        tx_buf = tx_buf[(MAX_BUF_SIZE // 2):] 

        # XBee -> FC
        if uart_xb.any():
            chunk = uart_xb.read()
            if chunk:
                rx_buf.extend(chunk)
            process_xbee_buffer()
            
        time.sleep_ms(1) 
        
    except MemoryError:
        # 發生 OOM 時的「斷尾求生」法，強制捨棄舊變數以釋放空間，防止死機
        tx_buf = bytearray()
        rx_buf = bytearray()
        gc.collect()
        time.sleep_ms(10)
        
    except Exception as e:
        time.sleep_ms(2)