// sketch.ino

#include "SystemConfig.h"
#include "Comms.h"
#include "Compression.h"
#include "Metrics.h"

// --- State Machine Definitions ---
enum SystemState {
    STATE_INIT,
    STATE_WAIT_INPUT,
    STATE_ANALYZE, // State added for logical clarity, though transition is fast
    STATE_TRANSMITTING,
    STATE_RECEIVING,
    STATE_REPORTING
};

SystemState currentState = STATE_INIT;
unsigned long stateStartTime = 0;

// Data Storage (Global state variables)
String rawDataBuffer = "";
String finalPayload = "";
int originalSize = 0;
int finalSize = 0;
String usedAlgorithm = "RAW";

// --- State Machine Functions ---

void Process_System_Logic(const String& rawData);
void TransitionTo(SystemState newState);

// --- SETUP ---
void setup() {
    HAL_UART_Init(BAUD_RATE);
    pinMode(TX_LED_PIN, OUTPUT);
    pinMode(RX_LED_PIN, OUTPUT);
    
    TransitionTo(STATE_INIT);
}

// --- MAIN LOOP ---
void loop() {
    // 1. Check for new serial input in all states
    if (Serial.available() > 0 && currentState == STATE_WAIT_INPUT) {
        rawDataBuffer = Serial.readStringUntil('\n');
        rawDataBuffer.trim();
        if (rawDataBuffer.length() > 0) {
            TransitionTo(STATE_ANALYZE); // <-- Transition to analyze state
        }
    }

    // 2. State-specific tasks (Non-blocking)
    switch (currentState) {
        case STATE_INIT:
            if (millis() > 100) { 
                SERIAL_PRINTLN_F("\n\n=================================================");
                SERIAL_PRINTLN_F("  UART INTELLIGENT COMPRESSION UNIT (V2.0 OPTIMIZED) ");
                SERIAL_PRINT_F("  TX/RX Delay: "); Serial.print(TX_RX_DELAY_MS); SERIAL_PRINTLN_F(" ms (Simulated)  ");
                SERIAL_PRINTLN_F("=================================================");
                TransitionTo(STATE_WAIT_INPUT);
            }
            break;
            
        case STATE_ANALYZE:
             // Execute the heavy logic when entering this state
             Process_System_Logic(rawDataBuffer); 
             // Logic inside Process_System_Logic will immediately call TransitionTo(STATE_TRANSMITTING);
             break;

        case STATE_TRANSMITTING:
            // CRITICAL CHECK: Non-blocking timer logic for TX simulation
            if (millis() - stateStartTime >= TX_RX_DELAY_MS) { 
                // Add a debug message here if it hangs again
                digitalWrite(TX_LED_PIN, LOW); 
                TransitionTo(STATE_RECEIVING);
            }
            break;

        case STATE_RECEIVING:
            // CRITICAL CHECK: Non-blocking timer logic for RX simulation
            if (millis() - stateStartTime >= TX_RX_DELAY_MS) {
                digitalWrite(RX_LED_PIN, LOW); 
                SERIAL_PRINTLN_F("[RX]: Verified! Generating Engineering Report...");
                TransitionTo(STATE_REPORTING);
            }
            break;
            
        case STATE_REPORTING:
            Print_Detailed_Metrics(rawDataBuffer, finalSize, usedAlgorithm);
            SERIAL_PRINTLN_F("-------------------------------------------------");
            SERIAL_PRINTLN_F("\n[SYSTEM]: Waiting for next input...");
            TransitionTo(STATE_WAIT_INPUT);
            break;
            
        default:
            // STATE_WAIT_INPUT is the default idle state
            break;
    }
}

// --- CORE LOGIC (Process & Decision) ---

void Process_System_Logic(const String& rawData) {
    SERIAL_PRINTLN_F("\n-------------------------------------------------");
    SERIAL_PRINT_F("[INPUT]: "); Print_Variable_String(rawData); Serial.println();
    
    // 1. ANALYSIS
    originalSize = rawData.length();
    String rlePayload = RunLengthEncoder(rawData);
    int rleSize = rlePayload.length();
    
    int huffmanSizeBytes = EstimateHuffmanSize(rawData);
    
    // 2. DECISION (Choose the smallest payload size)
    usedAlgorithm = "RAW";
    finalPayload = rawData;
    finalSize = originalSize;

    if (rleSize < finalSize) {
        finalPayload = rlePayload;
        finalSize = rleSize;
        usedAlgorithm = "RLE";
    }
    
    if (huffmanSizeBytes < finalSize) {
        finalPayload = rawData; 
        finalSize = huffmanSizeBytes; 
        usedAlgorithm = "HUFFMAN (Conceptual)";
    } 
    
    SERIAL_PRINT_F("[DECISION]: Algorithm -> "); Print_Variable_String(usedAlgorithm); Serial.println();
    
    // Transition immediately to TX State Preparation
    // This is the LAST step of the processing logic
    TransitionTo(STATE_TRANSMITTING);
}

// --- UTILITY ---
void TransitionTo(SystemState newState) {
    currentState = newState;
    stateStartTime = millis(); // <-- CRITICAL: Time is captured right here
    
    // Immediate actions on state entry
    switch(newState) {
        case STATE_TRANSMITTING:
            digitalWrite(TX_LED_PIN, HIGH); 
            SERIAL_PRINT_F("\n[TX]: Sending Data Packet... (Simulated "); 
            Serial.print(TX_RX_DELAY_MS); 
            SERIAL_PRINTLN_F("ms Delay)");
            Simulate_UART_Packet(finalPayload);
            break;
            
        case STATE_RECEIVING:
            SERIAL_PRINT_F("\n[RX]: Packet Arrived. Verifying Checksum... (Simulated "); 
            Serial.print(TX_RX_DELAY_MS); 
            SERIAL_PRINTLN_F("ms Delay)");
            digitalWrite(RX_LED_PIN, HIGH); 
            break;
            
        default:
            break;
    }
}