"""
Unit tests for the NIDS live pipeline (capture-service).
Covers flow key derivation, PPI matrix construction/padding, and model
classification/preprocessing logic — all without requiring live network
traffic or the running Docker network.
"""

import numpy as np
from scapy.all import IP, TCP

import flow_aggregator as fa
import inference_engine as ie


# ---------------------------------------------------------------------------
# flow_aggregator.py
# ---------------------------------------------------------------------------

def test_get_flow_key_same_for_both_directions():
    """Same conversation (two directions) must yield the same flow key."""
    pkt_forward = IP(src="10.0.0.1", dst="10.0.0.2") / TCP(sport=5000, dport=443)
    pkt_backward = IP(src="10.0.0.2", dst="10.0.0.1") / TCP(sport=443, dport=5000)

    key_fwd, _ = fa.get_flow_key(pkt_forward)
    key_bwd, _ = fa.get_flow_key(pkt_backward)

    assert key_fwd == key_bwd


def test_get_flow_key_none_for_non_tcp():
    """A packet without a TCP layer must return None."""
    pkt = IP(src="10.0.0.1", dst="10.0.0.2")
    assert fa.get_flow_key(pkt) is None


def test_compute_ppi_matrix_padding():
    """A flow shorter than SEQ_LEN must be zero-padded for the remaining rows."""
    flow = fa.create_new_flow_state(timestamp=1000.0, forward=True)
    flow["ppi_sizes"] = [100, 200, 150]
    flow["ppi_directions"] = [1, -1, 1]
    flow["ppi_times"] = [1000.0, 1000.05, 1000.12]

    matrix = fa.compute_ppi_matrix(flow)

    assert matrix.shape == (fa.SEQ_LEN, 3)
    # First packet: IAT is always 0 (no preceding packet)
    assert matrix[0, 0] == 0.0
    assert matrix[0, 1] == 1
    assert matrix[0, 2] == 100
    # Remaining rows (3 onward, past the actual sequence) must stay zero (padding)
    assert np.all(matrix[3:] == 0.0)


def test_compute_ppi_matrix_iat_calculation():
    """Verify IAT is computed in milliseconds between consecutive packets."""
    flow = fa.create_new_flow_state(timestamp=0.0, forward=True)
    flow["ppi_sizes"] = [100, 200]
    flow["ppi_directions"] = [1, 1]
    flow["ppi_times"] = [0.0, 0.5]  # 0.5 second gap = 500 milliseconds

    matrix = fa.compute_ppi_matrix(flow)

    assert matrix[1, 0] == 500.0


# ---------------------------------------------------------------------------
# inference_engine.py
# ---------------------------------------------------------------------------

def test_classify_flow_returns_three_values():
    """classify_flow must return (decision, error, latency_ms) with correct types."""
    zero_matrix = np.zeros((ie.model.seq_len, 3), dtype=np.float64)
    decision, error, latency_ms = ie.classify_flow(zero_matrix)

    assert decision in ("BENIGN", "MALICIOUS")
    assert isinstance(error, float)
    assert isinstance(latency_ms, float)
    assert latency_ms >= 0


def test_classify_flow_decision_matches_threshold():
    """The decision must logically match comparing the error against the threshold."""
    zero_matrix = np.zeros((ie.model.seq_len, 3), dtype=np.float64)
    decision, error, _ = ie.classify_flow(zero_matrix)

    expected = "MALICIOUS" if error > ie.THRESHOLD else "BENIGN"
    assert decision == expected


def test_preprocess_ppi_matrix_shape_preserved():
    """Preprocessing must preserve the matrix shape."""
    raw = np.random.rand(ie.model.seq_len, 3) * 100
    processed = ie.preprocess_ppi_matrix(raw)

    assert processed.shape == raw.shape


def test_preprocess_ppi_matrix_does_not_mutate_input():
    """preprocess_ppi_matrix must operate on a copy, not mutate the original matrix in place."""
    raw = np.ones((ie.model.seq_len, 3), dtype=np.float64) * 5.0
    raw_copy = raw.copy()

    ie.preprocess_ppi_matrix(raw)

    assert np.array_equal(raw, raw_copy)
