import socket
import struct
import time
import unittest

from app.flow_tracker import FlowTracker, _parse_frame
from app.live import TrafficNormalizer


class _Pipeline:
    def __init__(self):
        self.rows = []

    def publish(self, event, persist=True):
        self.rows.append((event, persist))


def _ipv4_tcp_frame(src: str, sport: int, dst: str, dport: int, payload: bytes = b"") -> bytes:
    eth = b"\x00" * 12 + struct.pack("!H", 0x0800)
    total_len = 20 + 20 + len(payload)
    ip = struct.pack(
        "!BBHHHBBH4s4s",
        0x45,
        0,
        total_len,
        1,
        0,
        64,
        6,
        0,
        socket.inet_aton(src),
        socket.inet_aton(dst),
    )
    tcp = struct.pack("!HHLLBBHHH", sport, dport, 0, 0, 5 << 4, 0x10, 65535, 0, 0)
    return eth + ip + tcp + payload


class FlowTrackerTests(unittest.TestCase):
    def test_parses_ipv4_tcp_tuple(self):
        frame = _ipv4_tcp_frame("192.168.100.10", 51000, "1.1.1.1", 443)
        self.assertEqual(_parse_frame(frame), ("192.168.100.10", 51000, "1.1.1.1", 443, "TCP"))

    def test_reverse_packets_update_one_live_session_without_persistence(self):
        pipeline = _Pipeline()
        tracker = FlowTracker(
            TrafficNormalizer("192.168.100.0/24"),
            pipeline,  # type: ignore[arg-type]
            update_interval_seconds=0.25,
            max_flows=1000,
        )
        outbound = _ipv4_tcp_frame("192.168.100.10", 51000, "1.1.1.1", 443, b"hello")
        inbound = _ipv4_tcp_frame("1.1.1.1", 443, "192.168.100.10", 51000, b"world")

        tracker.observe(outbound)
        time.sleep(0.26)
        tracker.observe(inbound)

        self.assertEqual(tracker.status()["active_flows"], 1)
        self.assertEqual(len(pipeline.rows), 2)
        first, first_persist = pipeline.rows[0]
        second, second_persist = pipeline.rows[1]
        self.assertEqual(first["id"], second["id"])
        self.assertEqual(second["direction"], "outbound")
        self.assertEqual(second["app_proto"], "tls")
        self.assertGreater(second["bytes"], first["bytes"])
        self.assertFalse(first_persist)
        self.assertFalse(second_persist)

    def test_status_exposes_monotonic_directional_traffic_counters(self):
        pipeline = _Pipeline()
        tracker = FlowTracker(
            TrafficNormalizer("192.168.100.0/24"),
            pipeline,  # type: ignore[arg-type]
            max_flows=1000,
        )
        outbound = _ipv4_tcp_frame("192.168.100.10", 51000, "1.1.1.1", 443, b"out")
        inbound = _ipv4_tcp_frame("1.1.1.1", 443, "192.168.100.10", 51000, b"in")
        internal = _ipv4_tcp_frame("192.168.100.10", 51000, "192.168.100.20", 443, b"lan")
        external = _ipv4_tcp_frame("1.1.1.1", 51000, "8.8.8.8", 443, b"wan")

        for frame in (outbound, inbound, internal, external):
            tracker.observe(frame)

        traffic = tracker.status()["traffic_counters"]
        self.assertEqual(traffic["packets_total"], 4)
        self.assertEqual(traffic["packets_out"], 1)
        self.assertEqual(traffic["packets_in"], 1)
        self.assertEqual(traffic["packets_internal"], 1)
        self.assertEqual(traffic["packets_external"], 1)
        self.assertEqual(
            traffic["bytes_total"],
            traffic["bytes_out"] + traffic["bytes_in"] + traffic["bytes_internal"] + traffic["bytes_external"],
        )


if __name__ == "__main__":
    unittest.main()