#include <Arduino.h>
#include <stdint.h>

// --- Configuration and Constants ---

// Maximum allowed payload size. Packets exceeding this size will be dropped.
// Set to the maximum supported by the protocol (4096 bytes).
#define MAX_PAYLOAD_SIZE 4096 
#define SYNC_WORD 0x7E
#define HEADER_SIZE 6 // The size *after* the SyncWord (Src ID, Dst ID, Payload Len)
#define DEBUG_PORT Serial //set debug port as serial (output)

// Chip IDs (MUST match the client configuration)
#define CLIENT_ID_1 0x0001
#define CLIENT_ID_2 0x0002
#define CLIENT_ID_3 0x0003

// Baud Rates
#define BAUD_RATE_1 38400  // For Serial1 (Client 1)
#define BAUD_RATE_2 19200  // For Serial2 (Client 2)
#define BAUD_RATE_3 4800   // For Serial3 (Client 3)

// Console Rate Gauge Configuration
#define GAUGE_UPDATE_INTERVAL_MS 1000UL // Update the data rate every 1 second , unsigned long , for millis

// --- State Management Structures (C-style) ---

typedef enum { 
    SYNC,
    HEADER_CONTENT,
    PAYLOAD_CONSUME,
    PAYLOAD_FORWARD
} State;

typedef struct {
    HardwareSerial* port;
    const char* name;
    State currentState;
    
    // Header components: 6 bytes 
    uint8_t headerBuffer[HEADER_SIZE]; 
    size_t headerIdx; 
    
    // Extracted data
    uint16_t sourceId;
    uint16_t destinationId;
    uint16_t payloadSize;
    
    // Payload tracking
    uint16_t payloadBytesRemaining;
    
    // Destination port for forwarding
    HardwareSerial* destPort;
} PortState_t;

// --- Global Instances ---

// Serial1 (Client 1), Serial2 (Client 2), Serial3 (Client 3)
PortState_t ports[] = {
    {&Serial1, "S1", SYNC, {0}, 0, 0, 0, 0, 0, NULL},
    {&Serial2, "S2", SYNC, {0}, 0, 0, 0, 0, 0, NULL},
    {&Serial3, "S3", SYNC, {0}, 0, 0, 0, 0, 0, NULL}
};
const size_t NUM_PORTS = sizeof(ports) / sizeof(ports[0]);

// Tracks bytes for the console rate gauge
uint32_t totalBytesForwarded = 0; 
uint32_t lastGaugeUpdate = 0;


// --- Helper Functions ---

/**
 * @brief Converts two Big-Endian bytes into a little-endian uint16_t.
 */
uint16_t parseBigEndian(const uint8_t* data) {
    // Big-endian: MSB is first byte (data[0]), LSB is second byte (data[1])
    return (uint16_t)data[0] << 8 | (uint16_t)data[1];
}

/**
 * @brief Resets the FSM state for a given port to the initial SYNC state.
 */
void resetState(PortState_t* state) {
    state->currentState = SYNC;
    state->headerIdx = 0;
    state->sourceId = 0;
    state->destinationId = 0;
    state->payloadSize = 0;
    state->payloadBytesRemaining = 0;
    state->destPort = NULL;
}

/**
 * @brief Safely writes data to a serial port, checking buffer availability.
 * This avoids the Arduino serial write() bug that can cause RX data loss.
 */
void safeSerialWrite(HardwareSerial* port, const uint8_t* data, size_t len) {
    size_t written = 0;
    while (written < len) {
        // Wait until there's space in the TX buffer
        if (port->availableForWrite() > 0) {
            size_t toWrite = min(len - written, (size_t)port->availableForWrite());
            port->write(&data[written], toWrite);
            written += toWrite;
        }
    }
}

/**
 * @brief Safely writes a single byte to a serial port.
 */
void safeSerialWriteByte(HardwareSerial* port, uint8_t byte) {
    while (port->availableForWrite() == 0) {
        // Wait for space in TX buffer
    }
    port->write(byte);
}

/**
 * @brief Central Finite State Machine (FSM) to process one byte from a source port.
 */
void processByte(PortState_t* state, uint8_t currentByte) {
    
    switch (state->currentState) {
        
        // --- State: SYNC ---
        case SYNC:
            if (currentByte == SYNC_WORD) {
                state->currentState = HEADER_CONTENT;
                state->headerIdx = 0; 
            }
            break;

        // --- State: HEADER_CONTENT ---
        case HEADER_CONTENT: {
            // Store the byte and advance the index
            state->headerBuffer[state->headerIdx++] = currentByte;

            if (state->headerIdx < HEADER_SIZE) {
                // Header is incomplete
                break; 
            }
            
            // --- Header Complete (6 bytes read) ---
            
            // 1. Extract Source ID (starts at index 0)
            state->sourceId = parseBigEndian(&state->headerBuffer[0]);
            
            // 2. Extract Destination ID (starts at index 2)
            state->destinationId = parseBigEndian(&state->headerBuffer[2]);

            // 3. Extract Payload Size (starts at index 4)
            state->payloadSize = parseBigEndian(&state->headerBuffer[4]);

            // 4. Size Check (Check against MAX_PAYLOAD_SIZE)
            if (state->payloadSize > MAX_PAYLOAD_SIZE) {
                DEBUG_PORT.print(state->name);
                DEBUG_PORT.print(": DROP (Too Long) Size: ");
                DEBUG_PORT.println(state->payloadSize);
                
                // Go to consume the remaining payload
                state->payloadBytesRemaining = state->payloadSize;
                state->currentState = PAYLOAD_CONSUME;
                break;
            }

            // 5. Determine Destination Port
            HardwareSerial* destPort = NULL;
            
            if (state->destinationId == CLIENT_ID_1) destPort = &Serial1;
            else if (state->destinationId == CLIENT_ID_2) destPort = &Serial2;
            else if (state->destinationId == CLIENT_ID_3) destPort = &Serial3;
            
            // 6. Validation Check: Is destination valid AND not the source port?
            if (destPort != NULL && destPort != state->port) {
                // Valid destination - prepare to forward
                state->destPort = destPort;
                state->payloadBytesRemaining = state->payloadSize;
                
                // Write sync word and header
                safeSerialWriteByte(destPort, SYNC_WORD);
                safeSerialWrite(destPort, state->headerBuffer, HEADER_SIZE);
                
                // Update byte counter (sync + header)
                totalBytesForwarded += (HEADER_SIZE + 1);
                
                // Transition to payload forwarding state
                if (state->payloadSize > 0) {
                    state->currentState = PAYLOAD_FORWARD;
                } else {
                    // No payload, packet complete
                    DEBUG_PORT.print(state->name);
                    DEBUG_PORT.print(" -> ");
                    DEBUG_PORT.print(destPort == &Serial1 ? "S1" : destPort == &Serial2 ? "S2" : "S3");
                    DEBUG_PORT.print(" | Fwd (Size: ");
                    DEBUG_PORT.print(HEADER_SIZE + 1);
                    DEBUG_PORT.println(" B)");
                    resetState(state);
                }
            } else {
                // Packet dropped (Invalid ID, destination is self, etc.)
                DEBUG_PORT.print(state->name);
                DEBUG_PORT.print(": DROP (Invalid/Self ID) Dest: 0x");
                DEBUG_PORT.println(state->destinationId, HEX);
                
                // Go to consume the payload if present, otherwise reset immediately.
                state->payloadBytesRemaining = state->payloadSize;
                state->currentState = (state->payloadSize > 0) ? PAYLOAD_CONSUME : SYNC;
            }
            break;
        }

        // --- State: PAYLOAD_FORWARD ---
        case PAYLOAD_FORWARD:
            // Forward this payload byte to destination
            safeSerialWriteByte(state->destPort, currentByte);
            totalBytesForwarded++;
            state->payloadBytesRemaining--;
            
            if (state->payloadBytesRemaining == 0) {
                // Packet complete
                uint16_t totalSize = state->payloadSize + HEADER_SIZE + 1;
                DEBUG_PORT.print(state->name);
                DEBUG_PORT.print(" -> ");
                DEBUG_PORT.print(state->destPort == &Serial1 ? "S1" : state->destPort == &Serial2 ? "S2" : "S3");
                DEBUG_PORT.print(" | Fwd (Size: ");
                DEBUG_PORT.print(totalSize);
                DEBUG_PORT.println(" B)");
                resetState(state);
            }
            break;

        // --- State: PAYLOAD_CONSUME (drop bytes) ---
        case PAYLOAD_CONSUME:
            // Discard this byte
            state->payloadBytesRemaining--;
            
            if (state->payloadBytesRemaining == 0) {
                // Done consuming, look for next sync
                resetState(state);
            }
            break;
    }
}


// --- ARDUINO SETUP AND LOOP ---

void setup() {
    // 1. DEBUG PORT (Monitor)
    DEBUG_PORT.begin(115200); 
    DEBUG_PORT.println("--- UART Router FSM Initialized ---");
    // Wait for serial port to connect (for native USB boards)
    while (!DEBUG_PORT) {}

    // 2. CLIENT PORTS (Communication)
    ports[0].port->begin(BAUD_RATE_1); // Client 1 (Serial1) -> 38400
    ports[1].port->begin(BAUD_RATE_2); // Client 2 (Serial2) -> 19200
    ports[2].port->begin(BAUD_RATE_3); // Client 3 (Serial3) -> 4800
    
    // Log the configured rates
    DEBUG_PORT.print("S1 (ID 0x"); DEBUG_PORT.print(CLIENT_ID_1, HEX); DEBUG_PORT.print("): "); DEBUG_PORT.println(BAUD_RATE_1);
    DEBUG_PORT.print("S2 (ID 0x"); DEBUG_PORT.print(CLIENT_ID_2, HEX); DEBUG_PORT.print("): "); DEBUG_PORT.println(BAUD_RATE_2);
    DEBUG_PORT.print("S3 (ID 0x"); DEBUG_PORT.print(CLIENT_ID_3, HEX); DEBUG_PORT.print("): "); DEBUG_PORT.println(BAUD_RATE_3);
    
    DEBUG_PORT.println("--- Router Ready. ---");
    lastGaugeUpdate = millis(); // Initialize the timer
}

void loop() {
    size_t i;
    uint32_t currentMillis = millis();
    
    // --- 1. Router Logic (High Priority) ---
    // Process bytes from all ports in round-robin fashion
    for (i = 0; i < NUM_PORTS; ++i) {
        PortState_t* state = &ports[i];
        
        // Process one byte at a time per port for fairness
        if (state->port->available() > 0) {
            uint8_t byte = state->port->read();
            processByte(state, byte);
        }
    }

    // --- 2. Console Gauge Update Logic (Low Priority) ---
    if (currentMillis - lastGaugeUpdate >= GAUGE_UPDATE_INTERVAL_MS) {
        // Calculate BPS: (Bytes * 1000) / Milliseconds
        uint32_t interval = currentMillis - lastGaugeUpdate; 
        uint32_t bytesPerSecond = (totalBytesForwarded * 1000UL) / interval;

        DEBUG_PORT.print("--- Router Rate: ");
        DEBUG_PORT.print(bytesPerSecond);
        DEBUG_PORT.println(" Bps ---");
        
        // Reset counter and timer
        totalBytesForwarded = 0;
        lastGaugeUpdate = currentMillis;
    }
}