// ============================================================================
// test/test_threat_correlator/test_threat_correlator.cpp
//
// Unity tests for src/threat_correlator.{h,cpp} — the rule-based multi-
// model detection-fusion engine. This is core, safety-relevant logic (it
// decides when independent weak detections should be fused into a single
// higher-confidence composite threat report) and previously had ZERO test
// coverage of any kind, native or otherwise -- it unconditionally
// `#include <Arduino.h>`'d, so it could not even be COMPILED outside the
// on-device build. Fixed by applying the same #ifdef ARDUINO / native-stub
// pattern already used throughout this codebase (network_id.cpp,
// wifi_security_scan.cpp, mythos_bridge.cpp) -- millis() is test-
// overridable via STUB_MILLIS_PROVIDED so the correlation-window and
// dedup timing can be driven deterministically instead of relying on
// wall-clock time.
//
// Run with:
//   pio test -e native -f test_threat_correlator
// ============================================================================

#include <unity.h>
#include <cstring>
#include <cmath>

// millis stub — test controls time via g_stub_millis
#define STUB_MILLIS_PROVIDED
static uint32_t g_stub_millis = 1000;
static inline uint32_t millis() { return g_stub_millis; }

// ---------------------------------------------------------------------------
// Stubs for threat_correlator.cpp's two gateway-feedback dependencies.
// Test-controllable so apply_gateway_overlay()'s two code paths (spoof
// penalty, PEC ceiling overlay) can be exercised deterministically.
// ---------------------------------------------------------------------------
static uint8_t g_stub_spoof_score = 0;
uint8_t spoof_gate_score_for_track(uint32_t) { return g_stub_spoof_score; }

static bool    g_stub_pec_available = false;
static float   g_stub_pec_probs[16] = {0};
static uint8_t g_stub_pec_n = 0;
bool pec_feedback_get_latest(uint32_t, float* probs_out, uint8_t* n_out, uint32_t) {
    if (!g_stub_pec_available) return false;
    memcpy(probs_out, g_stub_pec_probs, sizeof(float) * g_stub_pec_n);
    *n_out = g_stub_pec_n;
    return true;
}

#include "../../src/threat_correlator.cpp"

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
static ModelInferenceResult make_result(uint8_t slot, uint8_t cls, float conf) {
    ModelInferenceResult r{};
    r.slot_id = slot;
    r.class_label = cls;
    r.confidence = conf;
    r.num_classes = 0;
    r.latency_us = 0;
    r.valid = true;
    return r;
}

void setUp(void) {
    g_stub_millis = 1000;
    g_stub_spoof_score = 0;
    g_stub_pec_available = false;
    g_stub_pec_n = 0;
    correlator_init();
}
void tearDown(void) {}

// ---------------------------------------------------------------------------
// correlator_feed() / buffer management
// ---------------------------------------------------------------------------
void test_feed_valid_detection_increments_buffer(void) {
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    TEST_ASSERT_EQUAL_UINT8(1, correlator_buffer_count());
}

void test_feed_rejects_invalid_result(void) {
    ModelInferenceResult r = make_result(0, ATTACK_CLASS_DEAUTH, 0.8f);
    r.valid = false;
    correlator_feed(r);
    TEST_ASSERT_EQUAL_UINT8(0, correlator_buffer_count());
}

void test_feed_rejects_low_confidence(void) {
    // Below CORRELATOR_MIN_CONFIDENCE (0.50) must never enter the buffer.
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.49f));
    TEST_ASSERT_EQUAL_UINT8(0, correlator_buffer_count());
}

void test_reset_clears_buffer(void) {
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    TEST_ASSERT_EQUAL_UINT8(1, correlator_buffer_count());
    correlator_reset();
    TEST_ASSERT_EQUAL_UINT8(0, correlator_buffer_count());
}

// ---------------------------------------------------------------------------
// Rule 1: deauth + evil_twin within 5s -> credential_harvest
// ---------------------------------------------------------------------------
void test_rule1_credential_harvest_fires_within_window(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    g_stub_millis = 2000;
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_EQUAL(COMPOSITE_CREDENTIAL_HARVEST, threats[0].threat_class);
    TEST_ASSERT_EQUAL_UINT8(2, threats[0].constituent_count);
    TEST_ASSERT_TRUE(threats[0].valid);
    // Combined confidence = geometric mean of the two constituent confidences.
    TEST_ASSERT_FLOAT_WITHIN(0.001f, sqrtf(0.8f * 0.9f), threats[0].combined_confidence);
}

void test_rule1_does_not_fire_on_deauth_alone(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(0, n);
}

void test_rule1_does_not_fire_outside_window(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    // 6s later -- outside the 5s credential-harvest window.
    g_stub_millis = 1000 + CORR_RULE_CRED_HARVEST_WINDOW_MS + 1000;
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(0, n);
}

// ---------------------------------------------------------------------------
// Rule 2: drone_detected + gps_spoof within 10s -> drone_hijack_attempt
// ---------------------------------------------------------------------------
void test_rule2_drone_hijack_fires_within_window(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, CLASS_OCUSYNC, CONFIDENCE_THRESHOLD + 0.1f));
    g_stub_millis = 5000;
    correlator_feed(make_result(1, ATTACK_CLASS_GPS_SPOOF, 0.9f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_EQUAL(COMPOSITE_DRONE_HIJACK, threats[0].threat_class);
}

void test_rule2_drone_class_below_threshold_does_not_fire(void) {
    // Rule 2 requires the drone detection to clear CONFIDENCE_THRESHOLD
    // (stricter than the generic CORRELATOR_MIN_CONFIDENCE buffer gate) --
    // a low-confidence drone hint must not be enough on its own to trigger
    // a hijack-attempt composite.
    g_stub_millis = 1000;
    correlator_feed(make_result(0, CLASS_OCUSYNC, CORRELATOR_MIN_CONFIDENCE + 0.01f));
    g_stub_millis = 2000;
    correlator_feed(make_result(1, ATTACK_CLASS_GPS_SPOOF, 0.9f));

    TEST_ASSERT_TRUE_MESSAGE(CORRELATOR_MIN_CONFIDENCE + 0.01f < CONFIDENCE_THRESHOLD,
        "test setup assumption: fed confidence must be below CONFIDENCE_THRESHOLD");

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(0, n);
}

// ---------------------------------------------------------------------------
// Rule 3: jamming + replay within 5s -> mesh_compromise_attempt
// ---------------------------------------------------------------------------
void test_rule3_mesh_compromise_fires_within_window(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_JAMMING, 0.7f));
    g_stub_millis = 3000;
    correlator_feed(make_result(1, ATTACK_CLASS_REPLAY, 0.85f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_EQUAL(COMPOSITE_MESH_COMPROMISE, threats[0].threat_class);
}

// ---------------------------------------------------------------------------
// Rule 4: c2_beacon + any_attack -> active_intrusion
// ---------------------------------------------------------------------------
void test_rule4_active_intrusion_fires_with_any_attack_class(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_C2_BEACON, 0.75f));
    g_stub_millis = 4000;
    correlator_feed(make_result(1, ATTACK_CLASS_PORT_SCAN, 0.6f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_EQUAL(COMPOSITE_ACTIVE_INTRUSION, threats[0].threat_class);
}

// ---------------------------------------------------------------------------
// Multiple simultaneous rules in one cycle
// ---------------------------------------------------------------------------
void test_multiple_rules_fire_in_same_cycle(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));
    correlator_feed(make_result(2, ATTACK_CLASS_JAMMING, 0.7f));
    correlator_feed(make_result(3, ATTACK_CLASS_REPLAY, 0.85f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(2, n);
}

// ---------------------------------------------------------------------------
// Dedup: same composite class not re-emitted within CORRELATOR_DEDUP_MS
// ---------------------------------------------------------------------------
void test_dedup_suppresses_reemission_within_window(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));
    CompositeThreat threats[4];
    TEST_ASSERT_EQUAL_INT(1, correlator_run(threats, 4));

    // Fresh matching detections just after, still inside CORRELATOR_DEDUP_MS
    // of the first emission -- must be suppressed.
    g_stub_millis = 1000 + CORRELATOR_DEDUP_MS - 500;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, correlator_run(threats, 4),
        "same composite class must not re-fire within CORRELATOR_DEDUP_MS");
}

void test_dedup_allows_reemission_after_window(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));
    CompositeThreat threats[4];
    TEST_ASSERT_EQUAL_INT(1, correlator_run(threats, 4));

    // Past CORRELATOR_DEDUP_MS -- a fresh pairing should be allowed to fire
    // again. (Also comfortably inside CORRELATOR_WINDOW_MS=30s so the old
    // entries wouldn't just be pruned-and-refed; feed genuinely new ones.)
    g_stub_millis = 1000 + CORRELATOR_DEDUP_MS + 500;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));
    TEST_ASSERT_EQUAL_INT(1, correlator_run(threats, 4));
}

// ---------------------------------------------------------------------------
// Expired-entry pruning (CORRELATOR_WINDOW_MS)
// ---------------------------------------------------------------------------
void test_expired_entries_are_pruned_and_do_not_match(void) {
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));

    // Advance well past CORRELATOR_WINDOW_MS (30s) before the evil_twin
    // arrives -- the stale deauth must not still be sitting there to pair
    // with it.
    g_stub_millis = 1000 + CORRELATOR_WINDOW_MS + 5000;
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(0, n);
}

// ---------------------------------------------------------------------------
// apply_gateway_overlay(): spoof penalty + PEC ceiling
// ---------------------------------------------------------------------------
void test_gateway_overlay_spoof_penalty_reduces_confidence(void) {
    g_stub_spoof_score = 128;   // ~50% spoof probability
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.8f));
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.9f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(1, n);
    float unpenalized = sqrtf(0.8f * 0.9f);
    TEST_ASSERT_TRUE_MESSAGE(threats[0].combined_confidence < unpenalized,
        "a nonzero spoof score must reduce combined_confidence below the raw geometric mean");
}

void test_gateway_overlay_pec_ceiling_raises_confidence(void) {
    g_stub_pec_available = true;
    g_stub_pec_n = 4;
    g_stub_pec_probs[(uint8_t)COMPOSITE_CREDENTIAL_HARVEST] = 0.99f;
    g_stub_millis = 1000;
    correlator_feed(make_result(0, ATTACK_CLASS_DEAUTH, 0.5f));
    correlator_feed(make_result(1, ATTACK_CLASS_EVIL_TWIN, 0.5f));

    CompositeThreat threats[4];
    int n = correlator_run(threats, 4);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_FLOAT_WITHIN(0.001f, 0.99f, threats[0].combined_confidence);
}

// ---------------------------------------------------------------------------
// correlator_build_event()
// ---------------------------------------------------------------------------
void test_build_event_packs_fields_correctly(void) {
    CompositeThreat t{};
    t.threat_class = COMPOSITE_DRONE_HIJACK;
    t.combined_confidence = 0.75f;
    t.constituent_count = 2;
    t.constituent_slots[0] = 3; t.constituent_slots[1] = 5;
    t.constituent_slots[2] = 0; t.constituent_slots[3] = 0;
    t.constituent_labels[0] = CLASS_OCUSYNC;
    t.constituent_labels[1] = ATTACK_CLASS_GPS_SPOOF;
    t.valid = true;

    DetectionEvent evt;
    correlator_build_event(t, 0xDEADBEEFu, evt);

    TEST_ASSERT_EQUAL_UINT8(MSG_TYPE_COMPOSITE, evt.msg_type);
    TEST_ASSERT_EQUAL_UINT32(0xDEADBEEFu, evt.sensor_id);
    TEST_ASSERT_EQUAL_UINT8((uint8_t)(COMPOSITE_CLASS_BASE + (uint8_t)COMPOSITE_DRONE_HIJACK),
                             evt.class_label);
    TEST_ASSERT_EQUAL_UINT8((uint8_t)(0.75f * 255.0f), evt.confidence);
    TEST_ASSERT_EQUAL_UINT8(2, evt.signature_hash[0]);
    TEST_ASSERT_EQUAL_UINT8(CLASS_OCUSYNC, evt.signature_hash[3]);
    TEST_ASSERT_EQUAL_UINT8(ATTACK_CLASS_GPS_SPOOF, evt.signature_hash[4]);
    TEST_ASSERT_EQUAL_UINT8(FLAG_PROMOTED, evt.flags);
}

int main(int, char**) {
    UNITY_BEGIN();

    RUN_TEST(test_feed_valid_detection_increments_buffer);
    RUN_TEST(test_feed_rejects_invalid_result);
    RUN_TEST(test_feed_rejects_low_confidence);
    RUN_TEST(test_reset_clears_buffer);

    RUN_TEST(test_rule1_credential_harvest_fires_within_window);
    RUN_TEST(test_rule1_does_not_fire_on_deauth_alone);
    RUN_TEST(test_rule1_does_not_fire_outside_window);

    RUN_TEST(test_rule2_drone_hijack_fires_within_window);
    RUN_TEST(test_rule2_drone_class_below_threshold_does_not_fire);

    RUN_TEST(test_rule3_mesh_compromise_fires_within_window);

    RUN_TEST(test_rule4_active_intrusion_fires_with_any_attack_class);

    RUN_TEST(test_multiple_rules_fire_in_same_cycle);

    RUN_TEST(test_dedup_suppresses_reemission_within_window);
    RUN_TEST(test_dedup_allows_reemission_after_window);

    RUN_TEST(test_expired_entries_are_pruned_and_do_not_match);

    RUN_TEST(test_gateway_overlay_spoof_penalty_reduces_confidence);
    RUN_TEST(test_gateway_overlay_pec_ceiling_raises_confidence);

    RUN_TEST(test_build_event_packs_fields_correctly);

    return UNITY_END();
}
