/*
 * Parallel IDS with Watchdog Timer for Suspension Detection
 * - Hamming Distance (payload anomalies)
 * - ErrIDS Timing (inter-arrival anomalies)
 * - Watchdog Timer (message absence detection)
 */

#include <Arduino.h>
#include "training_ranges.h"
#include "timing_thresholds.h"
#include "can_traffic.h"

// ============================================================================
// CONFIG
// ============================================================================
#define MAX_ID_SLOTS 128
#define MAX_PAYLOAD_LEN 8

// Watchdog configuration
#define WATCHDOG_MULTIPLIER 3.0  // Trigger if message missing for 3x nominal period

// ============================================================================
// STRUCTURES
// ============================================================================
struct can_frame {
    uint16_t can_id;
    uint8_t can_dlc;
    uint8_t data[8];
};

struct CANIDState {
    uint16_t can_id;
    uint8_t last_payload[MAX_PAYLOAD_LEN];
    uint8_t payload_len;
    uint8_t min_hamming;
    uint8_t max_hamming;
    uint32_t last_timestamp;
    uint32_t nominal_period_ms_x1000;
    int32_t timing_lower_x1000;
    int32_t timing_upper_x1000;
    bool has_timing_params;
    bool initialized;
    
    // Watchdog fields
    uint32_t watchdog_deadline_ms;  // Next expected arrival time
    bool watchdog_active;           // Whether this ID needs monitoring
};

// ============================================================================
// GLOBALS
// ============================================================================
CANIDState id_states[MAX_ID_SLOTS];

// Metrics
uint32_t true_positives = 0;
uint32_t false_negatives = 0;
uint32_t false_positives = 0;
uint32_t true_negatives = 0;
uint32_t collateral_fps = 0;

// Method-specific detection counts
uint32_t hamming_detections = 0;
uint32_t erids_detections = 0;
uint32_t watchdog_detections = 0;
uint32_t combined_detections = 0;

// Timing
uint32_t total_detection_time_us = 0;
uint32_t detection_count = 0;
uint32_t min_detection_time_us = 999999;
uint32_t max_detection_time_us = 0;

uint8_t prev_label[MAX_ID_SLOTS];

// ============================================================================
// HASH FUNCTION
// ============================================================================
static inline uint8_t hash_can_id(uint16_t can_id) {
    return can_id % MAX_ID_SLOTS;
}

CANIDState* get_id_state(uint16_t can_id) {
    uint8_t slot = hash_can_id(can_id);
    uint8_t start_slot = slot;
    
    while (true) {
        if (id_states[slot].can_id == 0 || id_states[slot].can_id == can_id) {
            id_states[slot].can_id = can_id;
            return &id_states[slot];
        }
        slot = (slot + 1) % MAX_ID_SLOTS;
        if (slot == start_slot) return &id_states[0];
    }
}

// ============================================================================
// HAMMING
// ============================================================================
const uint8_t BIT_COUNT[16] = {0,1,1,2,1,2,2,3,1,2,2,3,2,3,3,4};

static inline uint8_t popcount(uint8_t b) {
    return BIT_COUNT[b & 0x0F] + BIT_COUNT[b >> 4];
}

uint8_t hamming_distance(const uint8_t* a, const uint8_t* b, uint8_t la, uint8_t lb) {
    uint8_t d = 0;
    uint8_t m = max(la, lb);
    for (uint8_t i = 0; i < m; i++) {
        uint8_t x = (i < la ? a[i] : 0) ^ (i < lb ? b[i] : 0);
        d += popcount(x);
    }
    return d;
}

// ============================================================================
// WATCHDOG CHECK - runs BEFORE processing current message
// ============================================================================
bool check_watchdog_violations(uint32_t current_time_ms, uint8_t* violated_label) {
    bool violation_detected = false;
    
    // Check all active watchdogs
    for (uint8_t slot = 0; slot < MAX_ID_SLOTS; slot++) {
        CANIDState* s = &id_states[slot];
        
        // Skip empty slots or inactive watchdogs
        if (s->can_id == 0 || !s->watchdog_active) continue;
        
        // Check if deadline expired
        if (current_time_ms > s->watchdog_deadline_ms) {
            violation_detected = true;
            
            // Check if this was a suspension attack (prev_label[slot] == 1 means last msg was attack)
            // For suspension, we expect NO messages, so label should indicate "attack in progress"
            // This is tricky - we'll conservatively assume violation = potential suspension
            
            // Record the violation (in real system would trigger alert)
            watchdog_detections++;
            
            // Reset deadline to avoid repeated alerts for same suspension
            // Add another period to the deadline
            uint32_t period_ms = s->nominal_period_ms_x1000 / 1000;
            s->watchdog_deadline_ms = current_time_ms + (uint32_t)(period_ms * WATCHDOG_MULTIPLIER);
        }
    }
    
    return violation_detected;
}

// ============================================================================
// DETECTION
// ============================================================================
bool detect_intrusion(can_frame* frame, uint32_t dataset_ts, uint32_t* dt_us) {
    uint32_t t0 = micros();

    uint16_t id = frame->can_id & 0x7FF;
    CANIDState* s = get_id_state(id);
    bool alert = false;
    bool err_alert = false;
    bool ham_alert = false;

    // ErrIDS timing
    if (s->initialized && s->has_timing_params) {
        uint32_t delta = dataset_ts - s->last_timestamp;
        int32_t grad = (int32_t)(delta * 1000) - (int32_t)s->nominal_period_ms_x1000;
        if (grad < s->timing_lower_x1000 || grad > s->timing_upper_x1000) {
            alert = true;
            err_alert = true;
            erids_detections++;
        }
    }

    // Hamming distance
    if (s->initialized && s->payload_len > 0) {
        uint8_t d = hamming_distance(frame->data, s->last_payload, frame->can_dlc, s->payload_len);
        if (d < s->min_hamming || d > s->max_hamming) {
            alert = true;
            ham_alert = true;
            hamming_detections++;
        }
    }

    // Update watchdog deadline for this ID (message arrived, so reset deadline)
    if (s->watchdog_active && s->has_timing_params) {
        uint32_t period_ms = s->nominal_period_ms_x1000 / 1000;
        s->watchdog_deadline_ms = dataset_ts + (uint32_t)(period_ms * WATCHDOG_MULTIPLIER);
    }

    // Always update state
    memcpy(s->last_payload, frame->data, frame->can_dlc);
    s->payload_len = frame->can_dlc;
    s->last_timestamp = dataset_ts;
    s->initialized = true;

    if (alert) combined_detections++;

    *dt_us = micros() - t0;
    return alert;
}

// ============================================================================
// SETUP
// ============================================================================
void setup() {
    Serial.begin(115200);
    delay(1000);

    memset(id_states, 0, sizeof(id_states));
    memset(prev_label, 0, sizeof(prev_label));

    // Load Hamming ranges
    for (uint16_t i = 0; i < NUM_RANGES; i++) {
        TrainingRange r;
        memcpy_P(&r, &TRAINING_RANGES[i], sizeof(r));
        CANIDState* s = get_id_state(r.can_id);
        s->min_hamming = r.min_hamming;
        s->max_hamming = r.max_hamming;
    }

    // Load timing thresholds and activate watchdog for periodic IDs
    uint16_t timing_loaded = 0;
    for (uint16_t i = 0; i < NUM_TIMING_THRESHOLDS; i++) {
        TimingThreshold t;
        memcpy_P(&t, &TIMING_THRESHOLDS[i], sizeof(t));
        
        CANIDState* s = get_id_state(t.can_id);
        s->has_timing_params = true;
        s->nominal_period_ms_x1000 = t.nominal_ms_x1000;
        s->timing_lower_x1000 = t.lower_threshold_x1000;
        s->timing_upper_x1000 = t.upper_threshold_x1000;
        
        // Activate watchdog for this ID
        s->watchdog_active = true;
        s->watchdog_deadline_ms = 0;  // Will be set on first message
        
        timing_loaded++;
    }

    // Memory report
    Serial.println("\n=== MEMORY ===");
    Serial.print("RAM: ");
    Serial.print(sizeof(id_states) + sizeof(prev_label));
    Serial.print("B (slots:");
    Serial.print(MAX_ID_SLOTS);
    Serial.print(" x ");
    Serial.print(sizeof(CANIDState));
    Serial.println("B)");
    Serial.print("Flash: ranges=");
    Serial.print(NUM_RANGES * sizeof(TrainingRange));
    Serial.print("B thr=");
    Serial.print(NUM_TIMING_THRESHOLDS * sizeof(TimingThreshold));
    Serial.println("B");
    Serial.print("Watchdog active for ");
    Serial.print(timing_loaded);
    Serial.println(" IDs");
    
    Serial.println("START");
}

// ============================================================================
// REPLAY LOOP
// ============================================================================
uint32_t replay_idx = 0;
bool done = false;

void loop() {
    if (done) return;

    if (replay_idx >= NUM_CAN_MESSAGES) {
        // Final metrics
        uint32_t avg = total_detection_time_us / detection_count;

        Serial.println("\n=== RESULTS ===");
        Serial.print("TP:");
        Serial.print(true_positives);
        Serial.print(" FN:");
        Serial.print(false_negatives);
        Serial.print(" FPp:");
        Serial.print(false_positives);
        Serial.print(" FPc:");
        Serial.print(collateral_fps);
        Serial.print(" TN:");
        Serial.println(true_negatives);
        
        Serial.println("\n=== DETECTION METHODS ===");
        Serial.print("Hamming:");
        Serial.print(hamming_detections);
        Serial.print(" ErrIDS:");
        Serial.print(erids_detections);
        Serial.print(" Watchdog:");
        Serial.print(watchdog_detections);
        Serial.print(" Combined:");
        Serial.println(combined_detections);
        
        Serial.println("\n=== PERFORMANCE ===");
        Serial.print("Avg:");
        Serial.print(avg);
        Serial.print("us Min:");
        Serial.print(min_detection_time_us);
        Serial.print("us Max:");
        Serial.print(max_detection_time_us);
        Serial.println("us");
        Serial.print("Throughput:");
        Serial.print(1000000 / avg);
        Serial.println("msg/s");
        
        Serial.println("\n=== MEMORY ===");
        Serial.print("RAM:");
        Serial.print(sizeof(id_states) + sizeof(prev_label));
        Serial.print("B Per-slot:");
        Serial.print(sizeof(CANIDState));
        Serial.println("B");
        Serial.print("Flash:ranges=");
        Serial.print(NUM_RANGES * sizeof(TrainingRange));
        Serial.print("B thr=");
        Serial.print(NUM_TIMING_THRESHOLDS * sizeof(TimingThreshold));
        Serial.println("B");

        done = true;
        return;
    }

    // Load message
    uint32_t curr_time_ms = pgm_read_dword(&CAN_TRAFFIC[replay_idx].timestamp_ms);
    uint16_t id = pgm_read_word(&CAN_TRAFFIC[replay_idx].can_id);
    uint8_t dlc = pgm_read_byte(&CAN_TRAFFIC[replay_idx].dlc);
    uint8_t label = pgm_read_byte(&CAN_TRAFFIC[replay_idx].label);

    // Check watchdog violations BEFORE processing current message
    uint8_t violated_label = 0;
    check_watchdog_violations(curr_time_ms, &violated_label);

    can_frame f;
    f.can_id = id;
    f.can_dlc = dlc;
    for (uint8_t i = 0; i < 8; i++)
        f.data[i] = pgm_read_byte(&CAN_TRAFFIC[replay_idx].data[i]);

    uint32_t dt;
    bool detected = detect_intrusion(&f, curr_time_ms, &dt);

    if (label == 1) {
        // ATTACK
        detected ? true_positives++ : false_negatives++;
    } else {
        // BENIGN
        if (detected) {
            uint8_t idx = hash_can_id(id);
            if (prev_label[idx] == 1) {
                collateral_fps++;
            } else {
                false_positives++;
            }
        } else {
            true_negatives++;
        }
    }
    
    // Update prev_label
    uint8_t idx = hash_can_id(id);
    prev_label[idx] = label;

    detection_count++;
    total_detection_time_us += dt;
    if (dt < min_detection_time_us) min_detection_time_us = dt;
    if (dt > max_detection_time_us) max_detection_time_us = dt;
    replay_idx++;
    
    // Progress
    if (replay_idx % 100 == 0) {
        Serial.print(replay_idx);
        Serial.print(" TP:");
        Serial.print(true_positives);
        Serial.print(" FP:");
        Serial.print(false_positives);
        Serial.print(" WD:");
        Serial.println(watchdog_detections);
    }
}