#include "threat_correlator.h"
#include "spoof_gate_client.h"
#include "pec_feedback.h"

#include <cstring>
#include <cmath>

#ifdef ARDUINO
#include <Arduino.h>
#else
// Host-side stubs so this module (previously buildable ONLY on-device) can
// be exercised by a native Unity test -- same pattern already used by
// network_id.cpp / mythos_bridge.cpp / wifi_security_scan.cpp. millis() is
// test-overridable via STUB_MILLIS_PROVIDED so tests can control the
// correlation-window/dedup timing deterministically.
#include <cstdio>
#include <cstdarg>
#ifndef STUB_MILLIS_PROVIDED
static uint32_t stub_millis_val = 1000;
static inline uint32_t millis() { return stub_millis_val; }
#endif
struct NativeSerialStub {
    void printf(const char* fmt, ...) {
        va_list args;
        va_start(args, fmt);
        vprintf(fmt, args);
        va_end(args);
    }
};
static NativeSerialStub Serial;
#endif // ARDUINO

static const char* TAG = "Correlator";

// ---------------------------------------------------------------------------
// Apply gateway feedback (rid_spoof_gate + Q-09 PEC) to a composite.
// Multiplies confidence by (1 - spoof_prob); overlays PEC-cleaned class prob
// when fresher than 5 s. track_id proxy = first constituent slot_id until a
// real RID-derived track_id is plumbed end-to-end.
// ---------------------------------------------------------------------------
static void apply_gateway_overlay(CompositeThreat& t) {
    uint32_t track_id = (uint32_t)t.constituent_slots[0];

    // 1. Spoof penalty
    uint8_t spoof_q8 = spoof_gate_score_for_track(track_id);
    if (spoof_q8 > 0) {
        float spoof_factor = 1.0f - (float)spoof_q8 / 255.0f;
        t.combined_confidence *= spoof_factor;
    }

    // 2. PEC ceiling overlay
    float pec_probs[MODEL_MAX_OUTPUT_CLASSES];
    uint8_t pec_n = 0;
    if (pec_feedback_get_latest(track_id, pec_probs, &pec_n, /*max_age_ms=*/5000)) {
        uint8_t cls_idx = (uint8_t)t.threat_class;
        if (cls_idx < pec_n) {
            float pec_conf = pec_probs[cls_idx];
            if (pec_conf > t.combined_confidence) {
                t.combined_confidence = pec_conf;
            }
        }
    }
}

// Rolling circular detection buffer
static CorrelatorDetection s_buffer[CORRELATOR_BUFFER_SIZE];
static uint8_t  s_buf_head = 0;
static uint8_t  s_buf_count = 0;

// Track which composites we've already emitted this cycle to avoid dups
static uint32_t s_last_composite_ms[COMPOSITE_CLASS_COUNT];
// Separate "has this class ever actually been emitted" flag -- needed
// because 0 is both zero-init's default value AND a legitimate millis()
// timestamp (the first ~CORRELATOR_DEDUP_MS after boot). Without this,
// is_duplicate() below would treat "never emitted" as "emitted at t=0",
// which meant the first-ever real composite threat of a given class was
// silently suppressed as a false "duplicate" whenever it happened to
// occur within CORRELATOR_DEDUP_MS (10s) of device boot.
static bool s_composite_ever_emitted[COMPOSITE_CLASS_COUNT];

// ---------------------------------------------------------------------------
// Internal: add a detection to the circular buffer
// ---------------------------------------------------------------------------
static void buffer_push(const CorrelatorDetection& det) {
    s_buffer[s_buf_head] = det;
    s_buf_head = (s_buf_head + 1) % CORRELATOR_BUFFER_SIZE;
    if (s_buf_count < CORRELATOR_BUFFER_SIZE) {
        s_buf_count++;
    }
}

// ---------------------------------------------------------------------------
// Internal: prune expired entries (older than CORRELATOR_WINDOW_MS)
// ---------------------------------------------------------------------------
static void prune_expired() {
    uint32_t now = millis();
    uint32_t cutoff = (now > CORRELATOR_WINDOW_MS) ? (now - CORRELATOR_WINDOW_MS) : 0;

    // Mark expired entries as consumed so they won't match
    for (int i = 0; i < s_buf_count; i++) {
        int idx = (s_buf_head - s_buf_count + i + CORRELATOR_BUFFER_SIZE) % CORRELATOR_BUFFER_SIZE;
        if (s_buffer[idx].timestamp_ms < cutoff) {
            s_buffer[idx].consumed = true;
        }
    }
}

// ---------------------------------------------------------------------------
// Internal: search buffer for a detection matching criteria within time window
// ---------------------------------------------------------------------------
struct SearchCriteria {
    uint8_t slot_id;        // MODEL_SLOT_INVALID = any slot
    uint8_t class_label;    // 0xFF = any class
    uint32_t max_age_ms;    // maximum age relative to reference timestamp
    uint32_t ref_time_ms;   // reference timestamp
    float    min_confidence;
};

#define MODEL_SLOT_ANY 0xFF

static const CorrelatorDetection* find_detection(const SearchCriteria& c) {
    const CorrelatorDetection* best = nullptr;
    float best_conf = 0.0f;

    for (int i = 0; i < s_buf_count; i++) {
        int idx = (s_buf_head - s_buf_count + i + CORRELATOR_BUFFER_SIZE) % CORRELATOR_BUFFER_SIZE;
        const CorrelatorDetection& d = s_buffer[idx];

        if (d.consumed) continue;

        // Check slot
        if (c.slot_id != MODEL_SLOT_ANY && d.slot_id != c.slot_id) continue;

        // Check class
        if (c.class_label != 0xFF && d.class_label != c.class_label) continue;

        // Check confidence
        if (d.confidence < c.min_confidence) continue;

        // Check time window
        uint32_t age;
        if (d.timestamp_ms <= c.ref_time_ms) {
            age = c.ref_time_ms - d.timestamp_ms;
        } else {
            age = d.timestamp_ms - c.ref_time_ms;
        }
        if (age > c.max_age_ms) continue;

        // Pick highest confidence match
        if (d.confidence > best_conf) {
            best_conf = d.confidence;
            best = &d;
        }
    }
    return best;
}

// ---------------------------------------------------------------------------
// Internal: build a CompositeThreat from constituent detections
// ---------------------------------------------------------------------------
static void build_composite(CompositeThreat& out,
                            CompositeThreatClass cls,
                            const CorrelatorDetection* dets[],
                            int count) {
    memset(&out, 0, sizeof(out));
    out.threat_class = cls;
    out.timestamp_ms = millis();
    out.valid = true;
    out.constituent_count = (uint8_t)count;

    float conf_product = 1.0f;
    for (int i = 0; i < count && i < 4; i++) {
        out.constituent_slots[i] = dets[i]->slot_id;
        out.constituent_labels[i] = dets[i]->class_label;
        out.constituent_confs[i] = dets[i]->confidence;
        conf_product *= dets[i]->confidence;
    }

    // Combined confidence = geometric mean
    if (count > 0) {
        out.combined_confidence = powf(conf_product, 1.0f / (float)count);
    }
}

// ---------------------------------------------------------------------------
// Internal: check dedup — don't emit the same composite class more than
// once per CORRELATOR_DEDUP_MS
// ---------------------------------------------------------------------------
static bool is_duplicate(CompositeThreatClass cls) {
    uint8_t idx = (uint8_t)cls;
    if (idx >= COMPOSITE_CLASS_COUNT) return false;
    if (!s_composite_ever_emitted[idx]) return false;

    uint32_t now = millis();
    if ((now - s_last_composite_ms[idx]) < CORRELATOR_DEDUP_MS) {
        return true;
    }
    return false;
}

static void mark_emitted(CompositeThreatClass cls) {
    uint8_t idx = (uint8_t)cls;
    if (idx < COMPOSITE_CLASS_COUNT) {
        s_last_composite_ms[idx] = millis();
        s_composite_ever_emitted[idx] = true;
    }
}

// ===========================================================================
// Public API
// ===========================================================================

void correlator_init() {
    memset(s_buffer, 0, sizeof(s_buffer));
    memset(s_last_composite_ms, 0, sizeof(s_last_composite_ms));
    memset(s_composite_ever_emitted, 0, sizeof(s_composite_ever_emitted));
    s_buf_head = 0;
    s_buf_count = 0;
    Serial.printf("[%s] Initialized. Buffer=%d slots, window=%d ms\n",
                  TAG, CORRELATOR_BUFFER_SIZE, CORRELATOR_WINDOW_MS);
}

// ---------------------------------------------------------------------------
void correlator_feed(const ModelInferenceResult& result) {
    if (!result.valid) return;

    // Only buffer detections above the correlation minimum confidence
    if (result.confidence < CORRELATOR_MIN_CONFIDENCE) return;

    CorrelatorDetection det;
    det.slot_id = result.slot_id;
    det.class_label = result.class_label;
    det.confidence = result.confidence;
    det.timestamp_ms = millis();
    det.consumed = false;

    buffer_push(det);
}

// ---------------------------------------------------------------------------
int correlator_run(CompositeThreat* threats, int max_threats) {
    prune_expired();

    int found = 0;
    uint32_t now = millis();

    // ------------------------------------------------------------------
    // Rule 1: deauth + evil_twin within 5s → credential_harvest
    // ------------------------------------------------------------------
    if (found < max_threats && !is_duplicate(COMPOSITE_CREDENTIAL_HARVEST)) {
        SearchCriteria c1 = { MODEL_SLOT_ANY, ATTACK_CLASS_DEAUTH, 5000, now, CORRELATOR_MIN_CONFIDENCE };
        SearchCriteria c2 = { MODEL_SLOT_ANY, ATTACK_CLASS_EVIL_TWIN, 5000, now, CORRELATOR_MIN_CONFIDENCE };

        const CorrelatorDetection* d1 = find_detection(c1);
        const CorrelatorDetection* d2 = find_detection(c2);

        if (d1 && d2) {
            // Check they are within 5s of each other
            uint32_t delta = (d1->timestamp_ms > d2->timestamp_ms)
                ? (d1->timestamp_ms - d2->timestamp_ms)
                : (d2->timestamp_ms - d1->timestamp_ms);
            if (delta <= CORR_RULE_CRED_HARVEST_WINDOW_MS) {
                const CorrelatorDetection* dets[] = { d1, d2 };
                build_composite(threats[found], COMPOSITE_CREDENTIAL_HARVEST, dets, 2);
                apply_gateway_overlay(threats[found]);
                mark_emitted(COMPOSITE_CREDENTIAL_HARVEST);
                found++;
                Serial.printf("[%s] COMPOSITE: credential_harvest (deauth+evil_twin)\n", TAG);
            }
        }
    }

    // ------------------------------------------------------------------
    // Rule 2: drone_detected + gps_spoof within 10s → drone_hijack_attempt
    // ------------------------------------------------------------------
    if (found < max_threats && !is_duplicate(COMPOSITE_DRONE_HIJACK)) {
        // drone_detected comes from the drone classifier (triage model, slot 0 typically)
        // We look for any drone class (CLASS_OCUSYNC, CLASS_FPV_FHSS, CLASS_SKYDIO, CLASS_UNKNOWN_DRONE)
        const CorrelatorDetection* d_drone = nullptr;
        for (uint8_t cl = CLASS_OCUSYNC; cl <= CLASS_UNKNOWN_DRONE; cl++) {
            SearchCriteria c = { MODEL_SLOT_ANY, cl, 10000, now, CONFIDENCE_THRESHOLD };
            d_drone = find_detection(c);
            if (d_drone) break;
        }

        SearchCriteria c_gps = { MODEL_SLOT_ANY, ATTACK_CLASS_GPS_SPOOF, 10000, now, CORRELATOR_MIN_CONFIDENCE };
        const CorrelatorDetection* d_gps = find_detection(c_gps);

        if (d_drone && d_gps) {
            uint32_t delta = (d_drone->timestamp_ms > d_gps->timestamp_ms)
                ? (d_drone->timestamp_ms - d_gps->timestamp_ms)
                : (d_gps->timestamp_ms - d_drone->timestamp_ms);
            if (delta <= CORR_RULE_DRONE_HIJACK_WINDOW_MS) {
                const CorrelatorDetection* dets[] = { d_drone, d_gps };
                build_composite(threats[found], COMPOSITE_DRONE_HIJACK, dets, 2);
                apply_gateway_overlay(threats[found]);
                mark_emitted(COMPOSITE_DRONE_HIJACK);
                found++;
                Serial.printf("[%s] COMPOSITE: drone_hijack_attempt (drone+gps_spoof)\n", TAG);
            }
        }
    }

    // ------------------------------------------------------------------
    // Rule 3: jamming + replay within 5s → mesh_compromise_attempt
    // ------------------------------------------------------------------
    if (found < max_threats && !is_duplicate(COMPOSITE_MESH_COMPROMISE)) {
        SearchCriteria c1 = { MODEL_SLOT_ANY, ATTACK_CLASS_JAMMING, 5000, now, CORRELATOR_MIN_CONFIDENCE };
        SearchCriteria c2 = { MODEL_SLOT_ANY, ATTACK_CLASS_REPLAY, 5000, now, CORRELATOR_MIN_CONFIDENCE };

        const CorrelatorDetection* d1 = find_detection(c1);
        const CorrelatorDetection* d2 = find_detection(c2);

        if (d1 && d2) {
            uint32_t delta = (d1->timestamp_ms > d2->timestamp_ms)
                ? (d1->timestamp_ms - d2->timestamp_ms)
                : (d2->timestamp_ms - d1->timestamp_ms);
            if (delta <= CORR_RULE_MESH_COMPROMISE_WINDOW_MS) {
                const CorrelatorDetection* dets[] = { d1, d2 };
                build_composite(threats[found], COMPOSITE_MESH_COMPROMISE, dets, 2);
                apply_gateway_overlay(threats[found]);
                mark_emitted(COMPOSITE_MESH_COMPROMISE);
                found++;
                Serial.printf("[%s] COMPOSITE: mesh_compromise (jamming+replay)\n", TAG);
            }
        }
    }

    // ------------------------------------------------------------------
    // Rule 4: c2_beacon + any_attack → active_intrusion
    // ------------------------------------------------------------------
    if (found < max_threats && !is_duplicate(COMPOSITE_ACTIVE_INTRUSION)) {
        SearchCriteria c_c2 = { MODEL_SLOT_ANY, ATTACK_CLASS_C2_BEACON, 10000, now, CORRELATOR_MIN_CONFIDENCE };
        const CorrelatorDetection* d_c2 = find_detection(c_c2);

        if (d_c2) {
            // Look for any attack-class detection
            const CorrelatorDetection* d_attack = nullptr;
            uint8_t attack_classes[] = {
                ATTACK_CLASS_DEAUTH, ATTACK_CLASS_EVIL_TWIN,
                ATTACK_CLASS_JAMMING, ATTACK_CLASS_REPLAY,
                ATTACK_CLASS_PORT_SCAN, ATTACK_CLASS_GPS_SPOOF
            };

            for (int a = 0; a < 6; a++) {
                SearchCriteria c = { MODEL_SLOT_ANY, attack_classes[a], 10000, now, CORRELATOR_MIN_CONFIDENCE };
                d_attack = find_detection(c);
                if (d_attack && d_attack != d_c2) break;
                d_attack = nullptr;
            }

            if (d_attack) {
                const CorrelatorDetection* dets[] = { d_c2, d_attack };
                build_composite(threats[found], COMPOSITE_ACTIVE_INTRUSION, dets, 2);
                apply_gateway_overlay(threats[found]);
                mark_emitted(COMPOSITE_ACTIVE_INTRUSION);
                found++;
                Serial.printf("[%s] COMPOSITE: active_intrusion (c2_beacon+attack)\n", TAG);
            }
        }
    }

    return found;
}

// ---------------------------------------------------------------------------
uint8_t correlator_buffer_count() {
    return s_buf_count;
}

// ---------------------------------------------------------------------------
void correlator_reset() {
    memset(s_buffer, 0, sizeof(s_buffer));
    s_buf_head = 0;
    s_buf_count = 0;
}

// ---------------------------------------------------------------------------
void correlator_build_event(const CompositeThreat& threat,
                            uint32_t sensor_id,
                            DetectionEvent& evt_out) {
    memset(&evt_out, 0, sizeof(evt_out));

    evt_out.msg_type = MSG_TYPE_COMPOSITE;
    evt_out.sensor_id = sensor_id;
    evt_out.timestamp_utc = (uint32_t)(millis() / 1000);

    // Map composite threat class to the class_label field
    // Composite classes start at COMPOSITE_CLASS_BASE
    evt_out.class_label = COMPOSITE_CLASS_BASE + (uint8_t)threat.threat_class;

    evt_out.confidence = (uint8_t)(threat.combined_confidence * 255.0f);

    // Pack constituent info into signature_hash:
    // [0] constituent_count
    // [1-2] constituent_slots (packed 2 per byte)
    // [3-6] constituent_labels
    // [7] reserved
    evt_out.signature_hash[0] = threat.constituent_count;
    evt_out.signature_hash[1] = (threat.constituent_slots[0] << 4) | (threat.constituent_slots[1] & 0x0F);
    evt_out.signature_hash[2] = (threat.constituent_slots[2] << 4) | (threat.constituent_slots[3] & 0x0F);
    for (int i = 0; i < 4; i++) {
        evt_out.signature_hash[3 + i] = threat.constituent_labels[i];
    }
    evt_out.signature_hash[7] = 0;

    evt_out.rssi = 0;
    evt_out.noise_floor = 0;
    evt_out.lat_fixed = 0;
    evt_out.lon_fixed = 0;
    evt_out.battery_pct = 0;  // caller can fill this in
    evt_out.flags = FLAG_PROMOTED; // composites are auto-promoted (skip FFPS)
}
