aivpn-server 0.4.0

//! Mask Generator — Analysis, Profile Building, and Self-Testing
//!
//! Analyzes recorded traffic metadata to generate MaskProfile:
//! 1. Statistical analysis (size distribution, IAT patterns, FSM states)
//! 2. Build MaskProfile from analysis results
//! 3. Self-test via Kolmogorov-Smirnov test
//! 4. Store and broadcast

// Generated masks carry all-zero ed25519 signatures until a dedicated
// signing key is plumbed through. Reject the build in production-secure mode
// to prevent unsigned masks from reaching end users.
#[cfg(feature = "production-secure")]
compile_error!(
    "mask_gen produces MaskProfile with signature=[0u8;64]. \
    Wire up a real Ed25519 signing key before enabling production-secure."
);

use std::sync::Arc;

use tracing::{error, info};

use aivpn_common::error::{Error, Result};
use aivpn_common::mask::*;
use aivpn_common::recording::{Direction, PacketMetadata};

use crate::mask_store::{MaskEntry, MaskStats, MaskStore};

// ─── Analysis Result ─────────────────────────────────────────────────────────

/// Result of traffic analysis
#[allow(dead_code)]
struct AnalysisResult {
    uplink: DirectionalAnalysis,
    downlink: DirectionalAnalysis,
    header: HeaderObservation,
    fsm_states: Vec<FSMState>,
    fsm_initial_state: u16,
    mean_entropy: f32,
    total_packets: u64,
    duration_secs: u64,
    confidence: f32,
}

#[allow(dead_code)]
struct DirectionalAnalysis {
    size_modes: Vec<Mode>,
    size_mean: f32,
    size_std: f32,
    iat_mean_ms: f32,
    iat_std_ms: f32,
    periods: Vec<Period>,
    packet_count: usize,
    /// Raw sorted sizes for empirical quantile-based distribution building
    raw_sizes_sorted: Vec<u16>,
    /// Raw sorted IATs for empirical distribution building
    raw_iats_sorted: Vec<f64>,
}

#[allow(dead_code)]
struct HeaderObservation {
    template: Vec<u8>,
    randomize_indices: Vec<usize>,
    header_spec: Option<HeaderSpec>,
    match_rate: f32,
    spec_confidence: f32,
}

/// Statistical mode (peak in distribution)
#[allow(dead_code)]
struct Mode {
    center: f32,
    std_dev: f32,
    weight: f32,
}

/// Periodic IAT pattern
#[allow(dead_code)]
struct Period {
    period_ms: f32,
    jitter_ms: f32,
    weight: f32,
}

/// Self-test result
struct SelfTestResult {
    ks_uplink_size: f32,
    ks_uplink_iat: f32,
    ks_downlink_size: f32,
    ks_downlink_iat: f32,
    header_match: f32,
    fsm_score: f32,
    entropy_penalty: f32,
    passed: bool,
    confidence: f32,
}

fn self_test_passes(
    ks_uplink_size: f32,
    ks_uplink_iat: f32,
    ks_downlink_size: f32,
    ks_downlink_iat: f32,
    header_match: f32,
    fsm_score: f32,
    entropy_penalty: f32,
    downlink_required: bool,
) -> bool {
    // With a correct KS implementation and empirical distributions, KS values
    // should be small (0.02-0.15).  Use generous thresholds to avoid rejecting
    // valid recordings while still catching bad profiles.
    let ks_ok = ks_uplink_size < 0.45
        && ks_uplink_iat < 0.45
        && (!downlink_required || (ks_downlink_size < 0.45 && ks_downlink_iat < 0.45));

    let structural_ok = header_match >= 0.55 && fsm_score >= 0.40 && entropy_penalty < 0.5;

    ks_ok && structural_ok
}

const MIN_ENCRYPTED_ENTROPY: f64 = 6.0;

// ─── Main Pipeline ───────────────────────────────────────────────────────────

/// Generate mask from recorded traffic and store it
pub async fn generate_and_store_mask(
    service: &str,
    packets: &[PacketMetadata],
    store: &Arc<MaskStore>,
) -> Result<String> {
    // 1. Analyze traffic
    let analysis = analyze_traffic(service, packets)?;
    info!(
        "Analysis complete for '{}': {} packets, up={} down={}, confidence={:.2}",
        service,
        analysis.total_packets,
        analysis.uplink.packet_count,
        analysis.downlink.packet_count,
        analysis.confidence
    );

    // 2. Build MaskProfile
    let profile = build_mask_profile(service, &analysis)?;

    // 3. Self-test
    let test = self_test(&profile, packets)?;
    if !test.passed {
        return Err(Error::Mask(format!(
            "Self-test failed: up(size={:.3},iat={:.3}) down(size={:.3},iat={:.3}) header={:.3} fsm={:.3} entropy_penalty={:.3}",
            test.ks_uplink_size,
            test.ks_uplink_iat,
            test.ks_downlink_size,
            test.ks_downlink_iat,
            test.header_match,
            test.fsm_score,
            test.entropy_penalty
        )));
    }
    info!(
        "Self-test passed for '{}': up(size={:.3},iat={:.3}) down(size={:.3},iat={:.3}) header={:.2} fsm={:.2}, confidence={:.2}",
        service,
        test.ks_uplink_size,
        test.ks_uplink_iat,
        test.ks_downlink_size,
        test.ks_downlink_iat,
        test.header_match,
        test.fsm_score,
        test.confidence
    );

    // 4. Store
    let mask_id = profile.mask_id.clone();
    store.add_mask(MaskEntry {
        profile,
        stats: MaskStats {
            mask_id: mask_id.clone(),
            times_used: 0,
            times_failed: 0,
            success_rate: 1.0,
            confidence: test.confidence,
            is_active: true,
            created_by: "auto".into(),
            created_at: current_unix_secs(),
            last_used: None,
        },
    })?;

    // 5. Broadcast to clients
    if let Err(e) = store.broadcast_mask_update(&mask_id).await {
        error!("Failed to broadcast mask '{}': {}", mask_id, e);
    }

    Ok(mask_id)
}

// ─── Traffic Analysis ────────────────────────────────────────────────────────

fn analyze_traffic(_service: &str, packets: &[PacketMetadata]) -> Result<AnalysisResult> {
    let uplink: Vec<&PacketMetadata> = packets
        .iter()
        .filter(|p| p.direction == Direction::Uplink)
        .collect();
    let downlink: Vec<&PacketMetadata> = packets
        .iter()
        .filter(|p| p.direction == Direction::Downlink)
        .collect();

    if uplink.len() < 100 {
        return Err(Error::Mask("Too few uplink packets (need >= 100)".into()));
    }

    let uplink_analysis = analyze_direction(&uplink);
    let downlink_analysis = analyze_direction(&downlink);

    // Header consensus
    let headers: Vec<Vec<u8>> = packets.iter().map(|p| p.header_prefix.clone()).collect();
    let header = analyze_headers(&headers);

    // FSM from size change-point detection
    let uplink_sizes: Vec<u16> = uplink.iter().map(|p| p.size).collect();
    let (fsm_states, fsm_initial) = build_fsm_from_sizes(&uplink_sizes);

    // Entropy
    let entropies: Vec<f32> = packets.iter().map(|p| p.entropy).collect();
    let mean_entropy = mean_f32(&entropies);

    // Confidence score
    let confidence = compute_confidence(
        packets.len(),
        uplink_analysis.packet_count,
        downlink_analysis.packet_count,
        header.match_rate,
        header.spec_confidence,
        mean_entropy,
    );

    // Duration
    let duration_secs = if packets.len() >= 2 {
        let first_ts = packets.first().map(|p| p.timestamp_ns).unwrap_or(0);
        let last_ts = packets.last().map(|p| p.timestamp_ns).unwrap_or(0);
        (last_ts.saturating_sub(first_ts)) / 1_000_000_000
    } else {
        0
    };

    Ok(AnalysisResult {
        uplink: uplink_analysis,
        downlink: downlink_analysis,
        header,
        fsm_states,
        fsm_initial_state: fsm_initial,
        mean_entropy,
        total_packets: packets.len() as u64,
        duration_secs,
        confidence,
    })
}

fn analyze_direction(packets: &[&PacketMetadata]) -> DirectionalAnalysis {
    let sizes: Vec<u16> = packets.iter().map(|p| p.size).collect();
    let iats: Vec<f64> = packets.iter().map(|p| p.iat_ms).collect();

    let mut raw_sizes_sorted = sizes.clone();
    raw_sizes_sorted.sort();
    let mut raw_iats_sorted = iats.clone();
    raw_iats_sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    DirectionalAnalysis {
        size_modes: find_modes_histogram(&sizes, 32),
        size_mean: mean_u16(&sizes),
        size_std: std_dev_u16(&sizes),
        iat_mean_ms: mean_f64(&iats) as f32,
        iat_std_ms: std_dev_f64(&iats) as f32,
        periods: find_periods(&iats),
        packet_count: packets.len(),
        raw_sizes_sorted,
        raw_iats_sorted,
    }
}

// ─── Mode Detection (Histogram) ─────────────────────────────────────────────

fn find_modes_histogram(sizes: &[u16], num_bins: usize) -> Vec<Mode> {
    if sizes.is_empty() {
        return vec![Mode {
            center: 64.0,
            std_dev: 32.0,
            weight: 1.0,
        }];
    }

    let min = *sizes.iter().min().unwrap_or(&0);
    let max = *sizes.iter().max().unwrap_or(&1500);
    let bin_width = ((max - min) as f32 / num_bins as f32).max(1.0);

    let mut bins = vec![0usize; num_bins];
    for &size in sizes {
        let bin = ((size as f32 - min as f32) / bin_width).min(num_bins as f32 - 1.0) as usize;
        bins[bin] += 1;
    }

    let total = sizes.len() as f32;
    let mut modes = Vec::new();

    for i in 1..bins.len().saturating_sub(1) {
        if bins[i] > bins[i - 1] && bins[i] > bins[i + 1] && bins[i] > total as usize / 20 {
            let center = min as f32 + (i as f32 + 0.5) * bin_width;
            let weight = bins[i] as f32 / total;

            // Compute local std dev around this mode
            let mut sum_sq = 0.0f32;
            let mut count = 0usize;
            for (j, &bin_count) in bins.iter().enumerate() {
                if bin_count > 0 {
                    let bc = min as f32 + (j as f32 + 0.5) * bin_width;
                    sum_sq += (bin_count as f32) * (bc - center).powi(2);
                    count += bin_count;
                }
            }
            let std_dev = if count > 0 {
                (sum_sq / count as f32).sqrt()
            } else {
                bin_width
            };
            modes.push(Mode {
                center,
                std_dev,
                weight,
            });
        }
    }

    // Fallback: single mode from mean/std
    if modes.is_empty() {
        let mean = mean_u16(sizes);
        let std = std_dev_u16(sizes);
        modes.push(Mode {
            center: mean,
            std_dev: std,
            weight: 1.0,
        });
    }

    modes
}

// ─── Period Detection (IAT) ─────────────────────────────────────────────────

fn find_periods(iats: &[f64]) -> Vec<Period> {
    if iats.len() < 10 {
        return vec![];
    }

    let mean_iat = mean_f64(iats);
    let std_iat = std_dev_f64(iats);

    if mean_iat < 1e-9 {
        return vec![];
    }

    if std_iat / mean_iat < 0.3 {
        // Stable single period
        vec![Period {
            period_ms: mean_iat as f32,
            jitter_ms: std_iat as f32,
            weight: 1.0,
        }]
    } else {
        // Bimodal — split by median
        let mut sorted = iats.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let median = sorted[sorted.len() / 2];

        let low: Vec<f64> = iats.iter().filter(|&&x| x <= median).copied().collect();
        let high: Vec<f64> = iats.iter().filter(|&&x| x > median).copied().collect();

        let mut periods = Vec::new();
        if !low.is_empty() {
            periods.push(Period {
                period_ms: mean_f64(&low) as f32,
                jitter_ms: std_dev_f64(&low) as f32,
                weight: low.len() as f32 / iats.len() as f32,
            });
        }
        if !high.is_empty() {
            periods.push(Period {
                period_ms: mean_f64(&high) as f32,
                jitter_ms: std_dev_f64(&high) as f32,
                weight: high.len() as f32 / iats.len() as f32,
            });
        }
        periods
    }
}

// ─── Change Point Detection → FSM ───────────────────────────────────────────

fn build_fsm_from_sizes(sizes: &[u16]) -> (Vec<FSMState>, u16) {
    if sizes.len() < 50 {
        return (
            vec![FSMState {
                state_id: 0,
                transitions: vec![],
            }],
            0,
        );
    }

    // 1. Detect change points (mean shift > 2σ in window of 20)
    let window = 20;
    let _global_mean = mean_u16(sizes);
    let global_std = std_dev_u16(sizes);
    let threshold = 2.0 * global_std;

    let mut change_points = Vec::new();
    for i in window..sizes.len().saturating_sub(window) {
        let before = sizes[i - window..i].iter().map(|&x| x as f32).sum::<f32>() / window as f32;
        let after = sizes[i..i + window].iter().map(|&x| x as f32).sum::<f32>() / window as f32;
        if (before - after).abs() > threshold {
            if change_points.is_empty() || i - *change_points.last().unwrap() > 10 {
                change_points.push(i);
            }
        }
    }

    // 2. Create segments
    let mut segments: Vec<&[u16]> = Vec::new();
    let mut start = 0;
    for &cp in &change_points {
        if cp > start {
            segments.push(&sizes[start..cp]);
        }
        start = cp;
    }
    if start < sizes.len() {
        segments.push(&sizes[start..]);
    }

    // Limit number of segments to avoid explosion
    if segments.is_empty() {
        segments.push(sizes);
    }

    // 3. Cluster segments by mean (threshold 100 bytes)
    let seg_means: Vec<f32> = segments
        .iter()
        .map(|s| s.iter().map(|&x| x as f32).sum::<f32>() / s.len().max(1) as f32)
        .collect();

    let mut clusters: Vec<Vec<usize>> = Vec::new();
    for (i, &seg_mean) in seg_means.iter().enumerate() {
        let mut assigned = false;
        for cluster in &mut clusters {
            let cm: f32 = cluster.iter().map(|&j| seg_means[j]).sum::<f32>() / cluster.len() as f32;
            if (seg_mean - cm).abs() < 100.0 {
                cluster.push(i);
                assigned = true;
                break;
            }
        }
        if !assigned {
            clusters.push(vec![i]);
        }
    }

    // Limit to max 8 FSM states
    clusters.truncate(8);

    // 4. Build FSM states with transitions
    let mut transitions: Vec<Vec<(u16, u32)>> = vec![vec![]; clusters.len()];
    for i in 0..segments.len().saturating_sub(1) {
        let from = clusters.iter().position(|c| c.contains(&i)).unwrap_or(0) as u16;
        let to = clusters
            .iter()
            .position(|c| c.contains(&(i + 1)))
            .unwrap_or(0) as u16;
        if (from as usize) < clusters.len() && (to as usize) < clusters.len() {
            if let Some(e) = transitions[from as usize]
                .iter_mut()
                .find(|(s, _)| *s == to)
            {
                e.1 += 1;
            } else {
                transitions[from as usize].push((to, 1));
            }
        }
    }

    // 5. Convert to FSMState
    let fsm_states: Vec<FSMState> = clusters
        .iter()
        .enumerate()
        .map(|(i, _cluster)| {
            let total: u32 = transitions[i].iter().map(|(_, c)| c).sum();
            let trans: Vec<FSMTransition> = transitions[i]
                .iter()
                .map(|(next, count)| FSMTransition {
                    condition: TransitionCondition::Random(*count as f32 / total.max(1) as f32),
                    next_state: *next,
                    size_override: None,
                    iat_override: None,
                    padding_override: None,
                })
                .collect();

            FSMState {
                state_id: i as u16,
                transitions: trans,
            }
        })
        .collect();

    (fsm_states, 0)
}

// ─── Header Consensus ────────────────────────────────────────────────────────

fn header_consensus(headers: &[Vec<u8>]) -> Vec<u8> {
    if headers.is_empty() {
        return vec![0u8; 8];
    }

    let len = headers.iter().map(|h| h.len()).min().unwrap_or(8).min(16);
    let mut result = Vec::with_capacity(len);

    for i in 0..len {
        let mut counts = [0u32; 256];
        for h in headers {
            if i < h.len() {
                counts[h[i] as usize] += 1;
            }
        }
        let max_count = counts.iter().max().copied().unwrap_or(0);
        let max_byte = counts.iter().position(|&c| c == max_count).unwrap_or(0) as u8;

        result.push(max_byte);
    }
    result
}

fn analyze_headers(headers: &[Vec<u8>]) -> HeaderObservation {
    let template = header_consensus(headers);
    let randomize_indices = header_randomize_indices(headers);
    let (header_spec, spec_confidence) = infer_header_spec(headers);
    let match_rate = header_spec
        .as_ref()
        .map(|spec| header_match_rate(spec, headers))
        .unwrap_or_else(|| raw_prefix_match_rate(&template, &randomize_indices, headers));

    HeaderObservation {
        template,
        randomize_indices,
        header_spec,
        match_rate,
        spec_confidence,
    }
}

fn header_randomize_indices(headers: &[Vec<u8>]) -> Vec<usize> {
    if headers.is_empty() {
        return vec![];
    }
    let max_len = headers.iter().map(|h| h.len()).min().unwrap_or(0).min(16);
    let mut randomize = Vec::new();
    for i in 0..max_len {
        let mut counts = [0u32; 256];
        for h in headers {
            if i < h.len() {
                counts[h[i] as usize] += 1;
            }
        }
        let max_count = counts.iter().max().copied().unwrap_or(0);
        let ratio = max_count as f32 / headers.len() as f32;
        if ratio < 0.85 {
            randomize.push(i);
        }
    }
    randomize
}

// ─── HeaderSpec Inference (Issue #30 fix) ─────────────────────────────────────

/// Infer HeaderSpec from observed traffic patterns
///
/// Analyzes the consistency of header bytes across packets to determine
/// if the traffic matches known protocol patterns (STUN, QUIC, DNS, TLS)
/// and generates an appropriate HeaderSpec for dynamic per-packet generation.
fn infer_header_spec(headers: &[Vec<u8>]) -> (Option<HeaderSpec>, f32) {
    if headers.is_empty() {
        return (None, 0.0);
    }

    // Analyze byte consistency at each position
    let max_len = headers.iter().map(|h| h.len()).max().unwrap_or(0).min(20);
    if max_len < 4 {
        return (None, 0.0); // Not enough data
    }

    // Calculate consistency ratio for each byte position
    let mut consistency: Vec<f32> = Vec::with_capacity(max_len);
    for i in 0..max_len {
        let mut counts = [0u32; 256];
        for h in headers {
            if i < h.len() {
                counts[h[i] as usize] += 1;
            }
        }
        let max_count = counts.iter().max().copied().unwrap_or(0);
        let ratio = max_count as f32 / headers.len() as f32;
        consistency.push(ratio);
    }

    let candidates = [
        score_stun(headers, &consistency),
        score_quic(headers, &consistency),
        score_dns(headers, &consistency),
        score_tls(headers, &consistency),
    ];

    let mut ranked: Vec<(HeaderSpec, f32)> = candidates.into_iter().flatten().collect();
    ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    if let Some((spec, score)) = ranked.first().cloned() {
        let second_best = ranked.get(1).map(|(_, score)| *score).unwrap_or(0.0);
        if score >= 0.72 && score - second_best >= 0.08 {
            return (Some(spec), score);
        }
    }

    // Fallback: use RawPrefix with randomization for variable bytes
    if headers.len() >= 4 {
        let template = headers[0].clone();
        let randomize_indices: Vec<usize> = consistency
            .iter()
            .enumerate()
            .filter(|(_, &c)| c < 0.7) // Randomize positions with <70% consistency
            .map(|(i, _)| i)
            .collect();

        if !randomize_indices.is_empty() && randomize_indices.len() < template.len() {
            return (
                Some(HeaderSpec::RawPrefix {
                    prefix_hex: hex::encode(&template),
                    randomize_indices,
                }),
                0.55,
            );
        }
    }

    // No pattern detected
    (None, 0.0)
}

fn score_stun(headers: &[Vec<u8>], consistency: &[f32]) -> Option<(HeaderSpec, f32)> {
    if headers.is_empty()
        || headers.iter().filter(|h| h.len() >= 20).count() * 10 < headers.len() * 7
    {
        return None;
    }

    let type_ratio = headers
        .iter()
        .filter(|h| h.len() >= 2 && h[0..2] == [0x00, 0x01])
        .count() as f32
        / headers.len() as f32;
    let cookie_ratio = headers
        .iter()
        .filter(|h| h.len() >= 8 && h[4..8] == [0x21, 0x12, 0xA4, 0x42])
        .count() as f32
        / headers.len() as f32;
    let len_sane_ratio = headers
        .iter()
        .filter(|h| h.len() >= 4 && h[2] & 0b1100_0000 == 0)
        .count() as f32
        / headers.len() as f32;
    let penalty = if headers
        .iter()
        .any(|h| !h.is_empty() && (0xC0..=0xCF).contains(&h[0]))
    {
        0.15
    } else {
        0.0
    };

    let score = (0.35 * type_ratio
        + 0.35 * cookie_ratio
        + 0.15 * len_sane_ratio
        + 0.15 * consistency.get(4).copied().unwrap_or(0.0))
        - penalty;

    if score > 0.45 {
        Some((
            HeaderSpec::stun_binding_with_cookie(cookie_ratio > 0.5),
            score.clamp(0.0, 1.0),
        ))
    } else {
        None
    }
}

fn score_quic(headers: &[Vec<u8>], consistency: &[f32]) -> Option<(HeaderSpec, f32)> {
    if headers.is_empty() {
        return None;
    }

    let long_header_ratio = headers
        .iter()
        .filter(|h| !h.is_empty() && (0xC0..=0xCF).contains(&h[0]))
        .count() as f32
        / headers.len() as f32;
    let version_ratio = headers
        .iter()
        .filter(|h| h.len() >= 5 && h[1..5] == [0x00, 0x00, 0x00, 0x01])
        .count() as f32
        / headers.len() as f32;
    let dcid_len_ratio = headers
        .iter()
        .filter(|h| h.len() >= 6 && (8..=20).contains(&h[5]))
        .count() as f32
        / headers.len() as f32;
    let penalty = if headers
        .iter()
        .filter(|h| h.len() >= 8 && h[4..8] == [0x21, 0x12, 0xA4, 0x42])
        .count() as f32
        / headers.len() as f32
        > 0.2
    {
        0.2
    } else {
        0.0
    };
    let score = (0.35 * long_header_ratio
        + 0.35 * version_ratio
        + 0.20 * dcid_len_ratio
        + 0.10 * consistency.get(0).copied().unwrap_or(0.0))
        - penalty;

    if score > 0.45 {
        Some((
            HeaderSpec::quic_initial(0x00000001, 8),
            score.clamp(0.0, 1.0),
        ))
    } else {
        None
    }
}

fn score_dns(headers: &[Vec<u8>], consistency: &[f32]) -> Option<(HeaderSpec, f32)> {
    if headers.is_empty()
        || headers.iter().filter(|h| h.len() >= 12).count() * 10 < headers.len() * 7
    {
        return None;
    }

    let flags_ratio = headers
        .iter()
        .filter(|h| h.len() >= 4 && h[2..4] == [0x01, 0x00])
        .count() as f32
        / headers.len() as f32;
    let counts_ratio = headers
        .iter()
        .filter(|h| h.len() >= 12 && h[4..12] == [0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
        .count() as f32
        / headers.len() as f32;
    let txid_variability = 1.0
        - consistency
            .get(0)
            .copied()
            .unwrap_or(1.0)
            .min(consistency.get(1).copied().unwrap_or(1.0));
    let penalty = if headers
        .iter()
        .any(|h| !h.is_empty() && (0xC0..=0xCF).contains(&h[0]))
    {
        0.15
    } else {
        0.0
    };
    let score = (0.35 * flags_ratio
        + 0.35 * counts_ratio
        + 0.20 * txid_variability
        + 0.10 * consistency.get(2).copied().unwrap_or(0.0))
        - penalty;

    if score > 0.45 {
        Some((HeaderSpec::dns_query(0x0100), score.clamp(0.0, 1.0)))
    } else {
        None
    }
}

fn score_tls(headers: &[Vec<u8>], consistency: &[f32]) -> Option<(HeaderSpec, f32)> {
    if headers.is_empty()
        || headers.iter().filter(|h| h.len() >= 5).count() * 10 < headers.len() * 7
    {
        return None;
    }

    let content_ratio = headers
        .iter()
        .filter(|h| h.len() >= 1 && matches!(h[0], 0x14 | 0x15 | 0x16 | 0x17))
        .count() as f32
        / headers.len() as f32;
    let version_ratio = headers
        .iter()
        .filter(|h| h.len() >= 3 && matches!(h[1..3], [0x03, 0x01] | [0x03, 0x02] | [0x03, 0x03]))
        .count() as f32
        / headers.len() as f32;
    let len_variability = 1.0
        - consistency
            .get(3)
            .copied()
            .unwrap_or(1.0)
            .min(consistency.get(4).copied().unwrap_or(1.0));
    let penalty = if headers
        .iter()
        .filter(|h| h.len() >= 8 && h[4..8] == [0x21, 0x12, 0xA4, 0x42])
        .count() as f32
        / headers.len() as f32
        > 0.2
    {
        0.15
    } else {
        0.0
    };
    let score = (0.35 * content_ratio
        + 0.35 * version_ratio
        + 0.15 * len_variability
        + 0.15 * consistency.get(0).copied().unwrap_or(0.0))
        - penalty;

    if score > 0.45 {
        let content_type = if headers
            .iter()
            .filter(|h| !h.is_empty() && h[0] == 0x17)
            .count()
            >= headers
                .iter()
                .filter(|h| !h.is_empty() && h[0] == 0x16)
                .count()
        {
            0x17
        } else {
            0x16
        };
        Some((
            HeaderSpec::tls_record(content_type, 0x0303),
            score.clamp(0.0, 1.0),
        ))
    } else {
        None
    }
}

// ─── Confidence Scoring ─────────────────────────────────────────────────────

fn compute_confidence(
    total_packets: usize,
    uplink_packets: usize,
    downlink_packets: usize,
    header_match_rate: f32,
    spec_confidence: f32,
    mean_entropy: f32,
) -> f32 {
    let mut score = 0.0f32;
    let min_dir = uplink_packets.min(downlink_packets) as f32;
    let max_dir = uplink_packets.max(downlink_packets).max(1) as f32;
    let direction_balance = min_dir / max_dir;

    if total_packets >= 10_000 {
        score += 0.3;
    } else if total_packets >= 5_000 {
        score += 0.25;
    } else if total_packets >= 1_000 {
        score += 0.2;
    } else if total_packets >= 500 {
        score += 0.15;
    }

    score += 0.2 * direction_balance.min(1.0);
    score += 0.2 * header_match_rate.min(1.0);
    score += 0.15 * spec_confidence.min(1.0);

    if mean_entropy > 7.0 {
        score += 0.15;
    } else if mean_entropy > 6.0 {
        score += 0.1;
    }

    score.min(1.0)
}

// ─── MaskProfile Builder ─────────────────────────────────────────────────────

fn build_mask_profile(service: &str, analysis: &AnalysisResult) -> Result<MaskProfile> {
    let mask_id = format!("auto_{}_v1", service.replace(' ', "_").to_lowercase());
    let header_template = analysis
        .header
        .header_spec
        .as_ref()
        .map(HeaderSpec::generate_static)
        .unwrap_or_else(|| analysis.header.template.clone());

    let size_distribution = build_size_distribution(&analysis.uplink);
    let iat_distribution = build_iat_distribution(&analysis.uplink);

    // Determine eph_pub_offset based on header_spec or header_template
    let eph_pub_offset = if let Some(ref spec) = analysis.header.header_spec {
        spec.min_length() as u16
    } else {
        header_template.len().min(4) as u16
    };

    // Determine spoof_protocol based on header_spec
    let spoof_protocol = analysis
        .header
        .header_spec
        .as_ref()
        .map(header_spec_protocol)
        .unwrap_or(SpoofProtocol::QUIC);

    let reverse_profile = build_reverse_profile(&mask_id, analysis);

    Ok(MaskProfile {
        mask_id,
        version: 2, // Version 2 for HeaderSpec support
        created_at: current_unix_secs(),
        expires_at: current_unix_secs() + 365 * 24 * 3600, // 1 year
        spoof_protocol,
        header_template,
        eph_pub_offset,
        eph_pub_length: 32,
        size_distribution,
        iat_distribution,
        padding_strategy: PaddingStrategy::RandomUniform { min: 0, max: 64 },
        fsm_states: analysis.fsm_states.clone(),
        fsm_initial_state: analysis.fsm_initial_state,
        signature_vector: vec![0.0; 64], // TODO: generate from neural model
        reverse_profile,
        signature: [0u8; 64], // TODO: sign with Ed25519
        header_spec: analysis.header.header_spec.clone(),
    })
}

fn build_reverse_profile(mask_id: &str, analysis: &AnalysisResult) -> Option<Box<MaskProfile>> {
    if analysis.downlink.packet_count < 50 {
        return None;
    }

    let size_distribution = build_size_distribution(&analysis.downlink);
    let iat_distribution = build_iat_distribution(&analysis.downlink);

    Some(Box::new(MaskProfile {
        mask_id: format!("{}_reverse", mask_id),
        version: 2,
        created_at: current_unix_secs(),
        expires_at: current_unix_secs() + 365 * 24 * 3600,
        spoof_protocol: analysis
            .header
            .header_spec
            .as_ref()
            .map(header_spec_protocol)
            .unwrap_or(SpoofProtocol::QUIC),
        header_template: analysis.header.template.clone(),
        eph_pub_offset: analysis
            .header
            .header_spec
            .as_ref()
            .map(|spec| spec.min_length() as u16)
            .unwrap_or(analysis.header.template.len().min(4) as u16),
        eph_pub_length: 32,
        size_distribution,
        iat_distribution,
        padding_strategy: PaddingStrategy::RandomUniform { min: 0, max: 64 },
        fsm_states: vec![FSMState {
            state_id: 0,
            transitions: vec![],
        }],
        fsm_initial_state: 0,
        signature_vector: vec![0.0; 64],
        reverse_profile: None,
        signature: [0u8; 64],
        header_spec: analysis.header.header_spec.clone(),
    }))
}

fn build_size_distribution(direction: &DirectionalAnalysis) -> SizeDistribution {
    // Frequency-weighted point bins: one bin per unique size value with
    // weight = frequency / total.  Sampling picks a bin by weight and returns
    // exactly that size value (min == max), perfectly reproducing the real CDF.
    // This avoids the uniform-within-range artefact of quantile range bins
    // that fills gaps between discrete packet sizes (MTU, ack, etc.).
    let sorted = &direction.raw_sizes_sorted;
    if sorted.is_empty() {
        return SizeDistribution {
            dist_type: SizeDistType::Histogram,
            bins: vec![(64, 512, 1.0)],
            parametric_type: None,
            parametric_params: None,
        };
    }

    let total = sorted.len() as f32;
    let mut bins: Vec<(u16, u16, f32)> = Vec::new();

    // Count frequencies via run-length on sorted data
    let mut i = 0;
    while i < sorted.len() {
        let val = sorted[i];
        let mut count = 0usize;
        while i < sorted.len() && sorted[i] == val {
            count += 1;
            i += 1;
        }
        bins.push((val, val, count as f32 / total));
    }

    SizeDistribution {
        dist_type: SizeDistType::Histogram,
        bins,
        parametric_type: None,
        parametric_params: None,
    }
}

fn build_iat_distribution(direction: &DirectionalAnalysis) -> IATDistribution {
    // Empirical quantile-based distribution: sample N evenly-spaced quantile
    // values from the sorted real IATs.  The Empirical sampler picks uniformly
    // among these stored values, faithfully reproducing the recorded CDF.
    let sorted = &direction.raw_iats_sorted;
    if sorted.is_empty() || sorted.len() < 10 {
        return IATDistribution {
            dist_type: IATDistType::LogNormal,
            params: vec![
                (direction.iat_mean_ms.max(1.0)).ln() as f64,
                (direction.iat_std_ms / direction.iat_mean_ms.max(1.0)).max(0.1) as f64,
            ],
            jitter_range_ms: symmetric_jitter_range(10.0),
        };
    }

    let num_quantiles = 200usize.min(sorted.len());
    let mut quantile_values: Vec<f64> = Vec::with_capacity(num_quantiles);
    for i in 0..num_quantiles {
        let idx = (i as f64 / (num_quantiles - 1).max(1) as f64 * (sorted.len() - 1) as f64).round()
            as usize;
        quantile_values.push(sorted[idx.min(sorted.len() - 1)]);
    }

    // Tiny jitter around quantile values to avoid exact duplicates
    let jitter = (direction.iat_std_ms.max(0.1) as f64) * 0.02;

    IATDistribution {
        dist_type: IATDistType::Empirical,
        params: quantile_values,
        jitter_range_ms: symmetric_jitter_range(jitter),
    }
}

fn symmetric_jitter_range(amplitude_ms: f64) -> (f64, f64) {
    let amplitude_ms = amplitude_ms.max(0.0);
    (-amplitude_ms, amplitude_ms)
}

// ─── Self-Test (Kolmogorov-Smirnov) ─────────────────────────────────────────

fn self_test(profile: &MaskProfile, packets: &[PacketMetadata]) -> Result<SelfTestResult> {
    // Generate synthetic samples from the mask profile
    let mut rng = rand::thread_rng();
    let uplink_count = packets
        .iter()
        .filter(|p| p.direction == Direction::Uplink)
        .count()
        .min(5000);
    let downlink_count = packets
        .iter()
        .filter(|p| p.direction == Direction::Downlink)
        .count()
        .min(5000);

    let synthetic_uplink_sizes: Vec<f64> = (0..uplink_count)
        .map(|_| profile.size_distribution.sample(&mut rng) as f64)
        .collect();
    let synthetic_uplink_iats: Vec<f64> = (0..uplink_count)
        .map(|_| profile.iat_distribution.sample(&mut rng))
        .collect();
    let synthetic_downlink_sizes: Vec<f64> = if let Some(ref reverse) = profile.reverse_profile {
        (0..downlink_count)
            .map(|_| reverse.size_distribution.sample(&mut rng) as f64)
            .collect()
    } else {
        vec![]
    };
    let synthetic_downlink_iats: Vec<f64> = if let Some(ref reverse) = profile.reverse_profile {
        (0..downlink_count)
            .map(|_| reverse.iat_distribution.sample(&mut rng))
            .collect()
    } else {
        vec![]
    };

    // Real data
    let real_uplink_sizes: Vec<f64> = packets
        .iter()
        .filter(|p| p.direction == Direction::Uplink)
        .take(5000)
        .map(|p| p.size as f64)
        .collect();
    let real_uplink_iats: Vec<f64> = packets
        .iter()
        .filter(|p| p.direction == Direction::Uplink)
        .take(5000)
        .map(|p| p.iat_ms)
        .collect();
    let real_downlink_sizes: Vec<f64> = packets
        .iter()
        .filter(|p| p.direction == Direction::Downlink)
        .take(5000)
        .map(|p| p.size as f64)
        .collect();
    let real_downlink_iats: Vec<f64> = packets
        .iter()
        .filter(|p| p.direction == Direction::Downlink)
        .take(5000)
        .map(|p| p.iat_ms)
        .collect();

    // KS tests
    let ks_uplink_size = ks_test(&synthetic_uplink_sizes, &real_uplink_sizes);
    let ks_uplink_iat = ks_test(&synthetic_uplink_iats, &real_uplink_iats);
    let ks_downlink_size =
        if !real_downlink_sizes.is_empty() && !synthetic_downlink_sizes.is_empty() {
            ks_test(&synthetic_downlink_sizes, &real_downlink_sizes)
        } else {
            1.0
        };
    let ks_downlink_iat = if !real_downlink_iats.is_empty() && !synthetic_downlink_iats.is_empty() {
        ks_test(&synthetic_downlink_iats, &real_downlink_iats)
    } else {
        1.0
    };

    // Entropy match
    let real_entropy: f64 =
        packets.iter().map(|p| p.entropy as f64).sum::<f64>() / packets.len().max(1) as f64;
    // Encrypted traffic entropy is length-sensitive for short packets, so use a
    // minimum floor instead of forcing an exact match to 7.0.
    let entropy_penalty = entropy_penalty(real_entropy);
    let headers: Vec<Vec<u8>> = packets.iter().map(|p| p.header_prefix.clone()).collect();
    let header_match = profile
        .header_spec
        .as_ref()
        .map(|spec| header_match_rate(spec, &headers))
        .unwrap_or_else(|| raw_prefix_match_rate(&profile.header_template, &[], &headers));
    let fsm_score = fsm_plausibility_score(profile, packets);
    let uplink = uplink_count as f32;
    let downlink = downlink_count as f32;
    let direction_balance = uplink.min(downlink.max(1.0)) / uplink.max(downlink).max(1.0);

    let downlink_required = downlink_count >= 100;
    let ks_threshold = 0.4;
    let passed = self_test_passes(
        ks_uplink_size,
        ks_uplink_iat,
        ks_downlink_size,
        ks_downlink_iat,
        header_match,
        fsm_score,
        entropy_penalty,
        downlink_required,
    );
    let confidence = if passed {
        let downlink_quality = if downlink_required {
            1.0 - ((ks_downlink_size + ks_downlink_iat) / 2.0 / 0.45).min(1.0)
        } else {
            0.25
        };
        let uplink_quality = 1.0 - ((ks_uplink_size + ks_uplink_iat) / 2.0 / ks_threshold).min(1.0);
        (0.45 * uplink_quality
            + 0.20 * downlink_quality
            + 0.20 * header_match
            + 0.10 * fsm_score
            + 0.03 * direction_balance.min(1.0)
            + 0.02 * (1.0 - entropy_penalty))
            .max(0.1)
    } else {
        0.0
    };

    Ok(SelfTestResult {
        ks_uplink_size,
        ks_uplink_iat,
        ks_downlink_size,
        ks_downlink_iat,
        header_match,
        fsm_score,
        entropy_penalty,
        passed,
        confidence,
    })
}

fn fsm_plausibility_score(profile: &MaskProfile, packets: &[PacketMetadata]) -> f32 {
    if profile.fsm_states.is_empty() {
        return 0.0;
    }

    let uplink_sizes: Vec<u16> = packets
        .iter()
        .filter(|p| p.direction == Direction::Uplink)
        .map(|p| p.size)
        .collect();
    if uplink_sizes.len() < 50 {
        return 0.5;
    }

    let (observed_states, observed_initial) = build_fsm_from_sizes(&uplink_sizes);
    let observed_count = observed_states.len().max(1) as f32;
    let profile_count = profile.fsm_states.len().max(1) as f32;
    let state_count_match =
        1.0 - ((observed_count - profile_count).abs() / observed_count.max(profile_count));

    let initial_match = if profile.fsm_initial_state == observed_initial {
        1.0
    } else {
        0.5
    };
    let has_valid_initial = if profile
        .fsm_states
        .iter()
        .any(|s| s.state_id == profile.fsm_initial_state)
    {
        1.0
    } else {
        0.0
    };
    let transition_mass = profile
        .fsm_states
        .iter()
        .map(|state| {
            let sum: f32 = state
                .transitions
                .iter()
                .map(|t| match t.condition {
                    TransitionCondition::Random(p) => p.max(0.0),
                    _ => 0.25,
                })
                .sum();
            (1.0 - (sum - 1.0).abs()).clamp(0.0, 1.0)
        })
        .sum::<f32>()
        / profile.fsm_states.len() as f32;

    (0.40 * state_count_match.max(0.0)
        + 0.25 * initial_match
        + 0.20 * has_valid_initial
        + 0.15 * transition_mass)
        .clamp(0.0, 1.0)
}

fn header_match_rate(spec: &HeaderSpec, headers: &[Vec<u8>]) -> f32 {
    if headers.is_empty() {
        return 0.0;
    }
    let matched = headers
        .iter()
        .filter(|header| header_matches_spec(spec, header))
        .count();
    matched as f32 / headers.len() as f32
}

fn raw_prefix_match_rate(template: &[u8], randomize_indices: &[usize], headers: &[Vec<u8>]) -> f32 {
    if headers.is_empty() {
        return 0.0;
    }
    let matched = headers
        .iter()
        .filter(|header| {
            let len = template.len().min(header.len());
            (0..len).all(|idx| randomize_indices.contains(&idx) || header[idx] == template[idx])
        })
        .count();
    matched as f32 / headers.len() as f32
}

fn header_matches_spec(spec: &HeaderSpec, header: &[u8]) -> bool {
    let fields = spec.fields();
    let total_len: usize = fields
        .iter()
        .map(|field| match field {
            HeaderField::Fixed { bytes } => bytes.len(),
            HeaderField::Random { len }
            | HeaderField::Length { len, .. }
            | HeaderField::Id { len, .. }
            | HeaderField::CounterLike { len, .. } => *len,
        })
        .sum();
    if header.len() < total_len {
        return false;
    }

    let mut cursor = 0usize;
    for field in fields {
        match field {
            HeaderField::Fixed { bytes } => {
                if header[cursor..cursor + bytes.len()] != bytes {
                    return false;
                }
                cursor += bytes.len();
            }
            HeaderField::Random { len } => {
                cursor += len;
            }
            HeaderField::Length { len, .. } => {
                cursor += len;
            }
            HeaderField::Id { len, mode } => {
                if matches!(mode, IdFieldMode::Zero)
                    && header[cursor..cursor + len].iter().any(|&b| b != 0)
                {
                    return false;
                }
                cursor += len;
            }
            HeaderField::CounterLike { len, .. } => {
                cursor += len;
            }
        }
    }
    true
}

fn header_spec_protocol(spec: &HeaderSpec) -> SpoofProtocol {
    match spec {
        HeaderSpec::RawPrefix { .. } => SpoofProtocol::QUIC,
        HeaderSpec::Structured { fields } => {
            if fields.len() >= 4
                && matches!(fields.first(), Some(HeaderField::Fixed { bytes }) if bytes == &vec![0x00, 0x01])
            {
                SpoofProtocol::WebRTC_STUN
            } else if fields.len() >= 4
                && matches!(fields.first(), Some(HeaderField::Fixed { bytes }) if bytes == &vec![0xC0])
            {
                SpoofProtocol::QUIC
            } else if fields.len() >= 3
                && matches!(fields.get(1), Some(HeaderField::Fixed { bytes }) if bytes == &vec![0x01, 0x00])
            {
                SpoofProtocol::DNS_over_UDP
            } else if fields.len() >= 3
                && matches!(fields.first(), Some(HeaderField::Fixed { bytes }) if bytes == &vec![0x16] || bytes == &vec![0x17] || bytes == &vec![0x15] || bytes == &vec![0x14])
            {
                SpoofProtocol::HTTPS_H2
            } else {
                SpoofProtocol::QUIC
            }
        }
    }
}

fn entropy_penalty(real_entropy: f64) -> f32 {
    if real_entropy >= MIN_ENCRYPTED_ENTROPY {
        0.0
    } else {
        (MIN_ENCRYPTED_ENTROPY - real_entropy).min(1.0) as f32
    }
}

/// Two-sample Kolmogorov-Smirnov test statistic
///
/// Correct merge-walk that handles tied values: at each unique value,
/// advance BOTH pointers past all duplicates before comparing CDFs.
fn ks_test(sample1: &[f64], sample2: &[f64]) -> f32 {
    if sample1.is_empty() || sample2.is_empty() {
        return 1.0;
    }

    let mut s1 = sample1.to_vec();
    let mut s2 = sample2.to_vec();
    s1.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    s2.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let n1 = s1.len() as f64;
    let n2 = s2.len() as f64;
    let mut max_diff: f64 = 0.0;

    let mut i = 0usize;
    let mut j = 0usize;

    while i < s1.len() || j < s2.len() {
        let v1 = if i < s1.len() { s1[i] } else { f64::INFINITY };
        let v2 = if j < s2.len() { s2[j] } else { f64::INFINITY };
        let current = v1.min(v2);

        // Advance both pointers past all elements <= current
        while i < s1.len() && s1[i] <= current {
            i += 1;
        }
        while j < s2.len() && s2[j] <= current {
            j += 1;
        }

        let cdf1 = i as f64 / n1;
        let cdf2 = j as f64 / n2;
        max_diff = max_diff.max((cdf1 - cdf2).abs());
    }

    max_diff as f32
}

// ─── Statistical Helpers ─────────────────────────────────────────────────────

fn mean_u16(data: &[u16]) -> f32 {
    if data.is_empty() {
        return 0.0;
    }
    data.iter().map(|&x| x as f32).sum::<f32>() / data.len() as f32
}

fn std_dev_u16(data: &[u16]) -> f32 {
    if data.is_empty() {
        return 0.0;
    }
    let m = mean_u16(data);
    let variance = data.iter().map(|&x| (x as f32 - m).powi(2)).sum::<f32>() / data.len() as f32;
    variance.sqrt()
}

fn mean_f64(data: &[f64]) -> f64 {
    if data.is_empty() {
        return 0.0;
    }
    data.iter().sum::<f64>() / data.len() as f64
}

fn std_dev_f64(data: &[f64]) -> f64 {
    if data.is_empty() {
        return 0.0;
    }
    let m = mean_f64(data);
    let variance = data.iter().map(|x| (x - m).powi(2)).sum::<f64>() / data.len() as f64;
    variance.sqrt()
}

fn mean_f32(data: &[f32]) -> f32 {
    if data.is_empty() {
        return 0.0;
    }
    data.iter().sum::<f32>() / data.len() as f32
}

fn current_unix_secs() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

#[cfg(test)]
mod tests {
    use super::{
        build_iat_distribution, entropy_penalty, header_consensus, header_match_rate,
        infer_header_spec, ks_test, self_test_passes, DirectionalAnalysis, Period,
    };
    use aivpn_common::mask::{HeaderSpec, IATDistType};

    #[test]
    fn entropy_penalty_accepts_short_encrypted_packets() {
        assert_eq!(entropy_penalty(6.03), 0.0);
        assert_eq!(entropy_penalty(7.8), 0.0);
        assert!(entropy_penalty(5.2) > 0.5);
    }

    #[test]
    fn multi_period_iat_distribution_preserves_spread_without_positive_bias() {
        // Generate synthetic IATs from two periods to test empirical quantile builder
        let mut raw_iats: Vec<f64> = Vec::new();
        // 60% near 20ms
        for i in 0..240 {
            raw_iats.push(15.0 + (i as f64 % 10.0));
        }
        // 40% near 100ms
        for i in 0..160 {
            raw_iats.push(70.0 + (i as f64 % 60.0));
        }
        raw_iats.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let direction = DirectionalAnalysis {
            size_modes: vec![],
            size_mean: 900.0,
            size_std: 120.0,
            iat_mean_ms: 52.0,
            iat_std_ms: 28.0,
            periods: vec![
                Period {
                    period_ms: 20.0,
                    jitter_ms: 5.0,
                    weight: 0.6,
                },
                Period {
                    period_ms: 100.0,
                    jitter_ms: 30.0,
                    weight: 0.4,
                },
            ],
            packet_count: 400,
            raw_sizes_sorted: vec![],
            raw_iats_sorted: raw_iats,
        };

        let dist = build_iat_distribution(&direction);
        assert_eq!(dist.dist_type, IATDistType::Empirical);
        assert!(dist.params.iter().any(|&value| value < 20.0));
        assert!(dist.params.iter().any(|&value| value > 100.0));
        assert!(dist.jitter_range_ms.0 < 0.0);
        assert!(dist.jitter_range_ms.1 > 0.0);
    }

    #[test]
    fn self_test_accepts_good_empirical_profile() {
        // With correct KS + empirical distributions, values should be small
        assert!(self_test_passes(
            0.05, 0.12, 0.08, 0.10, 1.000, 1.000, 0.000, true,
        ));
    }

    #[test]
    fn self_test_accepts_moderate_ks_with_good_structure() {
        // Some noise in IAT but still within bounds
        assert!(self_test_passes(
            0.15, 0.35, 0.20, 0.30, 0.80, 0.90, 0.000, true,
        ));
    }

    #[test]
    fn self_test_accepts_weak_header_above_threshold() {
        assert!(self_test_passes(
            0.10, 0.20, 0.15, 0.25, 0.60, 0.80, 0.000, true,
        ));
    }

    #[test]
    fn self_test_rejects_high_ks() {
        // KS > 0.45 should be rejected
        assert!(!self_test_passes(
            0.50, 0.60, 0.70, 0.80, 1.000, 1.000, 0.000, true,
        ));
    }

    #[test]
    fn self_test_still_rejects_broad_statistical_mismatch() {
        assert!(!self_test_passes(
            0.310, 0.820, 0.610, 0.590, 1.000, 1.000, 0.000, true,
        ));
    }

    #[test]
    fn ks_test_identical_samples_returns_zero() {
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        assert_eq!(ks_test(&data, &data), 0.0);
    }

    #[test]
    fn ks_test_identical_constant_returns_zero() {
        // This was the critical bug: all-same values caused KS ≈ 1.0
        let data: Vec<f64> = vec![1346.0; 5000];
        assert_eq!(ks_test(&data, &data), 0.0);
    }

    #[test]
    fn ks_test_similar_distributions_small() {
        let s1: Vec<f64> = (0..1000).map(|i| (i as f64) * 1.5).collect();
        let s2: Vec<f64> = (0..1000).map(|i| (i as f64) * 1.5 + 0.1).collect();
        let ks = ks_test(&s1, &s2);
        assert!(
            ks < 0.05,
            "KS for near-identical distributions should be small, got {}",
            ks
        );
    }

    #[test]
    fn header_consensus_is_deterministic_for_variable_bytes() {
        let headers = vec![
            vec![0x16, 0x03, 0x03, 0x00, 0x10],
            vec![0x16, 0x03, 0x03, 0x00, 0x20],
            vec![0x16, 0x03, 0x03, 0x00, 0x30],
        ];
        let c1 = header_consensus(&headers);
        let c2 = header_consensus(&headers);
        assert_eq!(c1, c2);
        assert_eq!(&c1[0..4], &[0x16, 0x03, 0x03, 0x00]);
    }

    #[test]
    fn infer_header_spec_rejects_random_noise() {
        let headers = vec![
            vec![0x91, 0x44, 0xF2, 0x7A, 0x19, 0x88, 0x01, 0x33],
            vec![0x6E, 0x10, 0xA3, 0x54, 0xC8, 0x92, 0x17, 0x45],
            vec![0x2D, 0xFF, 0x73, 0x0A, 0xB1, 0x4E, 0xC0, 0x18],
            vec![0x57, 0x28, 0x99, 0x61, 0x0F, 0xCD, 0xA4, 0x72],
        ];
        let (spec, confidence) = infer_header_spec(&headers);
        match spec {
            Some(HeaderSpec::RawPrefix { .. }) | None => {}
            other => panic!("unexpected protocol inference for noise: {:?}", other),
        }
        assert!(confidence <= 0.55);
    }

    #[test]
    fn infer_header_spec_detects_stun_and_matches_headers() {
        let headers = vec![
            vec![
                0x00, 0x01, 0x00, 0x00, 0x21, 0x12, 0xA4, 0x42, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
                12,
            ],
            vec![
                0x00, 0x01, 0x00, 0x08, 0x21, 0x12, 0xA4, 0x42, 13, 14, 15, 16, 17, 18, 19, 20, 21,
                22, 23, 24,
            ],
            vec![
                0x00, 0x01, 0x00, 0x10, 0x21, 0x12, 0xA4, 0x42, 25, 26, 27, 28, 29, 30, 31, 32, 33,
                34, 35, 36,
            ],
        ];
        let (spec, confidence) = infer_header_spec(&headers);
        let spec = spec.expect("stun should be inferred");
        assert!(confidence > 0.6);
        assert!(matches!(spec, HeaderSpec::Structured { .. }));
        assert!(header_match_rate(&spec, &headers) > 0.9);
    }
}