#include <Arduino.h>
#include <STM32FreeRTOS.h>
#include <cstdint>

constexpr uint8_t LED_PIN = A5;

//====================================================
// Structures
//====================================================
struct SensorData {
    int16_t temperature;
};

enum class SystemState : uint8_t {
    NORMAL,
    WARNING,
    CRITICAL
};

struct StatusData {
    int16_t temperature;
    SystemState state;
};

//====================================================
// Queues & Mutex
//====================================================
QueueHandle_t dataQueue;
QueueHandle_t statusQueue;
SemaphoreHandle_t uartMutex;

//====================================================
// Heartbeat Structures
//====================================================
struct TaskHealth {
    TickType_t lastHeartbeat;
};

struct Heartbeat {
    TaskHealth sensor;
    TaskHealth anomaly;
    TaskHealth logger;
};

volatile Heartbeat heartbeat;

//====================================================
// Watchdog Alert Flags
//====================================================
bool sensorAlertPrinted = false;
bool anomalyAlertPrinted = false;
bool loggerAlertPrinted = false;

//====================================================
// Fault Injection Flags
//====================================================
volatile bool injectSensorFault = false;
constexpr uint8_t SENSOR_FAILURE_AFTER = 20;

//====================================================
// Utility Functions
//====================================================
const char *StateToString(SystemState state) {
    switch (state) {
    case SystemState::NORMAL:
        return "NORMAL";
    case SystemState::WARNING:
        return "WARNING";
    case SystemState::CRITICAL:
        return "CRITICAL";
    default:
        return "UNKNOWN";
    }
}

//====================================================
// Temperature Thresholds
//====================================================
constexpr int16_t NORMAL_TO_WARNING = 35;
constexpr int16_t WARNING_TO_NORMAL = 33;
constexpr int16_t WARNING_TO_CRITICAL = 40;
constexpr int16_t CRITICAL_TO_WARNING = 38;

//====================================================
// Timing Constants
//====================================================
constexpr uint32_t SENSOR_PERIOD_MS = 300;
constexpr uint32_t WATCHDOG_PERIOD_MS = 300;
constexpr uint32_t WATCHDOG_TIMEOUT_MS = 1000;
//constexpr uint32_t QUEUE_TIMEOUT_MS = 50;

//====================================================
// Task Prototypes
//====================================================
void SensorTask(void *pvParameters);
void AnomalyTask(void *pvParameters);
void LoggerTask(void *pvParameters);
void WatchdogTask(void *pvParameters);

//====================================================
// Stack Overflow Hook
//====================================================
extern "C" void vApplicationStackOverflowHook(TaskHandle_t xTask, char *pcTaskName) {
    (void)xTask;
    (void)pcTaskName;
    // Rapidly toggle LED to indicate fatal stack overflow
    while (1) {
        digitalWrite(LED_PIN, !digitalRead(LED_PIN));
        for (volatile uint32_t i = 0; i < 50000; i++);
    }
}

//====================================================
// Watchdog Utility
//====================================================
void PrintWatchdogAlert(const char *taskName, TickType_t elapsedTicks) {
    Serial.println("================================");
    Serial.println("WATCHDOG ALERT");
    Serial.print("Task      : ");
    Serial.println(taskName);

    Serial.print("Elapsed   : ");
    Serial.print(elapsedTicks * portTICK_PERIOD_MS);
    Serial.println(" ms");

    Serial.print("Timeout   : ");
    Serial.print(WATCHDOG_TIMEOUT_MS);
    Serial.println(" ms");

    Serial.println("================================");
}

//====================================================
// setup()
//====================================================
void setup() {
    pinMode(LED_PIN, OUTPUT);
    Serial.begin(115200);
    delay(100);

    Serial.println();
    Serial.println("========================================");
    Serial.println("      FreeRTOS Task Monitor");
    Serial.println("========================================");
    Serial.println("Board      : STM32L031K6");
    Serial.println("Tasks      : 4");
    Serial.println("Queues     : 2");
    Serial.println("Watchdog   : Enabled");
    Serial.println("LED Status : Heartbeat");
    Serial.println("========================================");
    Serial.println();
    Serial.println("Commands:");
    Serial.println("f -> Inject Sensor Task Failure");
    Serial.println();

    dataQueue = xQueueCreate(2, sizeof(SensorData));
    statusQueue = xQueueCreate(2, sizeof(StatusData));
    uartMutex = xSemaphoreCreateMutex();

    if (dataQueue == NULL || statusQueue == NULL || uartMutex == NULL) {
        Serial.println("Error: Failed to create Queues/Mutex (Out of RAM)");
        while (1);
    }

    // Stack sizes balanced for 8KB total RAM on STM32L031K6
    BaseType_t r1 = xTaskCreate(SensorTask,   "Sensor",   90,  NULL, 1, NULL);
    BaseType_t r2 = xTaskCreate(AnomalyTask,  "Anomaly",  90,  NULL, 1, NULL);
    BaseType_t r3 = xTaskCreate(LoggerTask,   "Logger",   140, NULL, 1, NULL);
    BaseType_t r4 = xTaskCreate(WatchdogTask, "Watchdog", 110, NULL, 2, NULL);

    if (r1 != pdPASS || r2 != pdPASS || r3 != pdPASS || r4 != pdPASS) {
        Serial.println("Error: Failed to create Tasks (Out of RAM)");
        while (1);
    }

    Serial.println("Starting Scheduler...");
    vTaskStartScheduler();
}

void loop() {
    // Idle state handled by FreeRTOS scheduler
}

//====================================================
// Sensor Task
//====================================================
void SensorTask(void *pvParameters) {
    (void)pvParameters;
    SensorData sensorData;
    int16_t temperature = 20;
    uint8_t cycleCount = 0;

    TickType_t lastWakeTime = xTaskGetTickCount();
    while (1) {
        sensorData.temperature = temperature;

        xQueueSend(dataQueue, &sensorData, portMAX_DELAY);
          heartbeat.sensor.lastHeartbeat = xTaskGetTickCount();

        temperature++;
        if (temperature > 40) {
            temperature = 20;
        }
        cycleCount++;

        // Fault injection logic
        if (injectSensorFault && cycleCount >= SENSOR_FAILURE_AFTER) {
            if (xSemaphoreTake(uartMutex, pdMS_TO_TICKS(100)) == pdTRUE) {
                Serial.println();
                Serial.println("Injecting Sensor Task Failure (Suspending Task)...");
                Serial.println();
                xSemaphoreGive(uartMutex);
            }
            // Suspend task without burning CPU cycles in a tight loop
            vTaskSuspend(NULL);
        }

        vTaskDelayUntil(&lastWakeTime, pdMS_TO_TICKS(SENSOR_PERIOD_MS));
    }
}

//====================================================
// Anomaly Task
//====================================================
void AnomalyTask(void *pvParameters) {
    (void)pvParameters;
    SensorData receivedData;
    StatusData statusData;
    SystemState currentState = SystemState::NORMAL;

    while (1) {
        if (xQueueReceive(dataQueue, &receivedData, portMAX_DELAY) == pdTRUE) {
            switch (currentState) {
            case SystemState::NORMAL:
                if (receivedData.temperature >= NORMAL_TO_WARNING) {
                    currentState = SystemState::WARNING;
                }
                break;
            case SystemState::WARNING:
                if (receivedData.temperature <= WARNING_TO_NORMAL) {
                    currentState = SystemState::NORMAL;
                } else if (receivedData.temperature >= WARNING_TO_CRITICAL) {
                    currentState = SystemState::CRITICAL;
                }
                break;
            case SystemState::CRITICAL:
                if (receivedData.temperature <= CRITICAL_TO_WARNING) {
                    currentState = SystemState::WARNING;
                }
                break;
            }

            statusData.temperature = receivedData.temperature;
            statusData.state = currentState;

            xQueueSend(statusQueue, &statusData, portMAX_DELAY);
              heartbeat.anomaly.lastHeartbeat = xTaskGetTickCount();
        }
    }
}

//====================================================
// Logger Task
//====================================================
void LoggerTask(void *pvParameters) {
    (void)pvParameters;
    StatusData receivedData;
    bool firstMessage = true;
    SystemState previousState = SystemState::NORMAL;

    while (1) {
        if (xQueueReceive(statusQueue, &receivedData, portMAX_DELAY) == pdTRUE) {
            if (xSemaphoreTake(uartMutex, portMAX_DELAY) == pdTRUE) {
                if (firstMessage) {
                    Serial.println("================================");
                    Serial.println("SYSTEM STARTED");
                    Serial.print("Current State : ");
                    Serial.println(StateToString(receivedData.state));
                    Serial.println("================================");
                    previousState = receivedData.state;
                    firstMessage = false;
                } else if (previousState != receivedData.state) {
                    Serial.println("================================");
                    Serial.println("STATE TRANSITION");
                    Serial.print(StateToString(previousState));
                    Serial.print(" -> ");
                    Serial.println(StateToString(receivedData.state));
                    Serial.println("================================");
                    previousState = receivedData.state;
                }

                Serial.print("Temp: ");
                Serial.print(receivedData.temperature);
                Serial.print(" C | State: ");
                Serial.println(StateToString(receivedData.state));

                xSemaphoreGive(uartMutex);
                heartbeat.logger.lastHeartbeat = xTaskGetTickCount();
            }
        }
    }
}

//====================================================
// Watchdog Task
//====================================================
void WatchdogTask(void *pvParameters) {
    (void)pvParameters;
    bool ledState = false;

    TickType_t initTick = xTaskGetTickCount();
    heartbeat.sensor.lastHeartbeat = initTick;
    heartbeat.anomaly.lastHeartbeat = initTick;
    heartbeat.logger.lastHeartbeat = initTick;

    while (1) {
        // Handle incoming Serial commands inside WatchdogTask to save RAM
        if (Serial.available()) {
            char cmd = Serial.read();
            if (cmd == 'f' || cmd == 'F') {
                injectSensorFault = true;
                if (xSemaphoreTake(uartMutex, pdMS_TO_TICKS(50)) == pdTRUE) {
                    Serial.println();
                    Serial.println("Fault Injection Enabled");
                    Serial.println();
                    xSemaphoreGive(uartMutex);
                }
            }
        }

        TickType_t currentTick = xTaskGetTickCount();

        ledState = !ledState;
        digitalWrite(LED_PIN, ledState);

        TickType_t elapsedSensor = currentTick - heartbeat.sensor.lastHeartbeat;
        TickType_t elapsedAnomaly = currentTick - heartbeat.anomaly.lastHeartbeat;
        TickType_t elapsedLogger = currentTick - heartbeat.logger.lastHeartbeat;

        if (xSemaphoreTake(uartMutex, pdMS_TO_TICKS(50)) == pdTRUE) {
            if (elapsedSensor > pdMS_TO_TICKS(WATCHDOG_TIMEOUT_MS) && !sensorAlertPrinted) {
                PrintWatchdogAlert("Sensor", elapsedSensor);
                sensorAlertPrinted = true;
            }

            if (elapsedAnomaly > pdMS_TO_TICKS(WATCHDOG_TIMEOUT_MS) && !anomalyAlertPrinted) {
                PrintWatchdogAlert("Anomaly", elapsedAnomaly);
                anomalyAlertPrinted = true;
            }

            if (elapsedLogger > pdMS_TO_TICKS(WATCHDOG_TIMEOUT_MS) && !loggerAlertPrinted) {
                PrintWatchdogAlert("Logger", elapsedLogger);
                loggerAlertPrinted = true;
            }

            xSemaphoreGive(uartMutex);
        }

        vTaskDelay(pdMS_TO_TICKS(WATCHDOG_PERIOD_MS));
    }
}