"""Network pcap analysis — read-only over stored captures.

Per FN-A0.1 decision #1: dpkt for analysis (fast, dependency-light);
tcpdump shell-out for the human-readable dump (CLI-side, FN-A8).

Every analyzer:

- Streams the pcap (no full-file load) so a 5 GB capture doesn't OOM.
- Returns a frozen dataclass + JSON serializer.
- Is a pure function of (pcap_path, optional config). Reanalysis is
  deterministic.

dpkt's `Reader` yields ``(timestamp, raw_bytes)`` per packet. We
parse Ethernet → IP/IPv6 → TCP/UDP/ICMP and pull out:
- src/dst IPs + L4 ports for top-talkers
- DNS queries (UDP port 53; A/AAAA/CNAME extraction)
- TLS SNI from the ClientHello
- HTTP method/path/host from plaintext TCP streams (request only)
"""

from __future__ import annotations

import logging
import struct
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator
    from pathlib import Path


__all__ = [
    "DnsQuery",
    "HttpRequest",
    "NetworkAnalysis",
    "ProtocolBreakdown",
    "TalkerPair",
    "TlsSni",
    "analyze_pcap",
    "iter_packets",
]


LOG = logging.getLogger("mgf.vm.forensics.analysis.network")


# ---------------------------------------------------------------------------
# Result dataclasses
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class TalkerPair:
    """One (src, dst) connection summary."""

    src: str
    dst: str
    bytes: int
    packets: int


@dataclass(frozen=True)
class ProtocolBreakdown:
    """Counts of packets per L3/L4 protocol."""

    tcp: int = 0
    udp: int = 0
    icmp: int = 0
    other_l4: int = 0
    ipv4: int = 0
    ipv6: int = 0
    arp: int = 0
    other_l3: int = 0


@dataclass(frozen=True)
class DnsQuery:
    """One DNS query observed in the capture."""

    name: str
    qtype: str
    first_seen_at: datetime


@dataclass(frozen=True)
class TlsSni:
    """One TLS SNI hostname observed."""

    hostname: str
    first_seen_at: datetime


@dataclass(frozen=True)
class HttpRequest:
    """One plaintext HTTP request observed."""

    method: str
    host: str
    path: str
    seen_at: datetime


@dataclass(frozen=True)
class NetworkAnalysis:
    """Bundle of every network-side analyzer output for one pcap."""

    pcap_path: str
    packet_count: int
    bytes_total: int
    duration_s: float
    started_at: datetime | None
    ended_at: datetime | None
    protocols: ProtocolBreakdown
    top_talkers: tuple[TalkerPair, ...]
    dns_queries: tuple[DnsQuery, ...]
    tls_snis: tuple[TlsSni, ...]
    http_requests: tuple[HttpRequest, ...]
    bpf_filter: str = ""
    errors: tuple[str, ...] = field(default_factory=tuple)


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


def analyze_pcap(
    pcap_path: Path,
    *,
    top_n_talkers: int = 20,
    sample_dns: int = 200,
    sample_sni: int = 200,
    sample_http: int = 200,
) -> NetworkAnalysis:
    """Stream-read ``pcap_path`` and return a :class:`NetworkAnalysis`.

    Caps on the per-list samples (``sample_dns`` etc.) keep memory
    bounded even for very chatty captures. Set to 0 to disable a
    list entirely.
    """
    talkers: dict[tuple[str, str], list[int]] = defaultdict(lambda: [0, 0])
    proto_counter: Counter[str] = Counter()
    dns_seen: dict[tuple[str, str], DnsQuery] = {}
    sni_seen: dict[str, TlsSni] = {}
    http_seen: list[HttpRequest] = []
    errors: list[str] = []

    packet_count = 0
    bytes_total = 0
    first_ts: float | None = None
    last_ts: float | None = None

    for ts, raw in iter_packets(pcap_path):
        packet_count += 1
        bytes_total += len(raw)
        first_ts = ts if first_ts is None else min(first_ts, ts)
        last_ts = ts if last_ts is None else max(last_ts, ts)

        try:
            _process_packet(
                ts=ts,
                raw=raw,
                talkers=talkers,
                proto_counter=proto_counter,
                dns_seen=dns_seen,
                sni_seen=sni_seen,
                http_seen=http_seen,
                sample_dns=sample_dns,
                sample_sni=sample_sni,
                sample_http=sample_http,
            )
        except Exception as e:
            # Per-packet parse failures must not abort the whole
            # analysis. Log + count + continue.
            errors.append(f"packet {packet_count}: {e}")
            if len(errors) > 50:
                # Bound the error list — the first 50 are usually
                # representative.
                continue

    # Derive top-N talkers.
    top_talkers = sorted(
        (
            TalkerPair(src=src, dst=dst, bytes=stats[0], packets=stats[1])
            for (src, dst), stats in talkers.items()
        ),
        key=lambda t: (-t.bytes, -t.packets),
    )[:top_n_talkers]

    started_at = datetime.fromtimestamp(first_ts, tz=UTC) if first_ts is not None else None
    ended_at = datetime.fromtimestamp(last_ts, tz=UTC) if last_ts is not None else None
    duration = (last_ts - first_ts) if (first_ts is not None and last_ts is not None) else 0.0

    return NetworkAnalysis(
        pcap_path=str(pcap_path),
        packet_count=packet_count,
        bytes_total=bytes_total,
        duration_s=duration,
        started_at=started_at,
        ended_at=ended_at,
        protocols=ProtocolBreakdown(
            tcp=proto_counter["tcp"],
            udp=proto_counter["udp"],
            icmp=proto_counter["icmp"],
            other_l4=proto_counter["other_l4"],
            ipv4=proto_counter["ipv4"],
            ipv6=proto_counter["ipv6"],
            arp=proto_counter["arp"],
            other_l3=proto_counter["other_l3"],
        ),
        top_talkers=tuple(top_talkers),
        dns_queries=tuple(sorted(dns_seen.values(), key=lambda d: d.first_seen_at)),
        tls_snis=tuple(sorted(sni_seen.values(), key=lambda s: s.first_seen_at)),
        http_requests=tuple(http_seen),
        errors=tuple(errors),
    )


# ---------------------------------------------------------------------------
# Per-packet dispatch
# ---------------------------------------------------------------------------


def _process_packet(
    *,
    ts: float,
    raw: bytes,
    talkers: dict[tuple[str, str], list[int]],
    proto_counter: Counter[str],
    dns_seen: dict[tuple[str, str], DnsQuery],
    sni_seen: dict[str, TlsSni],
    http_seen: list[HttpRequest],
    sample_dns: int,
    sample_sni: int,
    sample_http: int,
) -> None:
    import dpkt
    import dpkt.dns as dpkt_dns
    import dpkt.ethernet as dpkt_eth
    import dpkt.ip as dpkt_ip
    import dpkt.ip6 as dpkt_ip6

    eth = dpkt_eth.Ethernet(raw)

    # ARP early-out.
    if isinstance(eth.data, dpkt.arp.ARP):
        proto_counter["arp"] += 1
        return

    if isinstance(eth.data, dpkt_ip.IP):
        proto_counter["ipv4"] += 1
        ip = eth.data
        src = _ipv4_str(ip.src)
        dst = _ipv4_str(ip.dst)
    elif isinstance(eth.data, dpkt_ip6.IP6):
        proto_counter["ipv6"] += 1
        ip = eth.data
        src = _ipv6_str(ip.src)
        dst = _ipv6_str(ip.dst)
    else:
        proto_counter["other_l3"] += 1
        return

    # Track per-pair byte + packet counts.
    key = (src, dst)
    talkers[key][0] += len(raw)
    talkers[key][1] += 1

    # L4 dispatch.
    transport = ip.data
    if isinstance(transport, dpkt.tcp.TCP):
        proto_counter["tcp"] += 1
        _maybe_extract_http(transport, src, dst, ts, http_seen, sample_http)
        _maybe_extract_sni(transport, ts, sni_seen, sample_sni)
    elif isinstance(transport, dpkt.udp.UDP):
        proto_counter["udp"] += 1
        if transport.sport == 53 or transport.dport == 53:
            _maybe_extract_dns(
                transport,
                ts,
                dns_seen,
                sample_dns,
                dpkt_dns,
            )
    elif isinstance(transport, (dpkt.icmp.ICMP, dpkt.icmp6.ICMP6)):
        proto_counter["icmp"] += 1
    else:
        proto_counter["other_l4"] += 1


# ---------------------------------------------------------------------------
# DNS extraction
# ---------------------------------------------------------------------------


def _maybe_extract_dns(
    udp: Any,
    ts: float,
    dns_seen: dict[tuple[str, str], DnsQuery],
    sample_cap: int,
    dpkt_dns: Any,
) -> None:
    if sample_cap == 0:
        return  # 0 = disabled
    if len(dns_seen) >= sample_cap:
        return
    try:
        msg = dpkt_dns.DNS(udp.data)
    except Exception:
        return
    # Process queries (qr=0). Responses are useful too but for v1
    # we just record the question section per query.
    for q in msg.qd:
        name = q.name if isinstance(q.name, str) else q.name.decode("ascii", errors="replace")
        qtype = _dns_qtype_name(q.type)
        key = (name, qtype)
        if key not in dns_seen:
            dns_seen[key] = DnsQuery(
                name=name,
                qtype=qtype,
                first_seen_at=datetime.fromtimestamp(ts, tz=UTC),
            )


_DNS_QTYPE_NAMES: dict[int, str] = {
    1: "A",
    2: "NS",
    5: "CNAME",
    6: "SOA",
    12: "PTR",
    15: "MX",
    16: "TXT",
    28: "AAAA",
    33: "SRV",
    65: "HTTPS",
}


def _dns_qtype_name(code: int) -> str:
    return _DNS_QTYPE_NAMES.get(code, f"TYPE{code}")


# ---------------------------------------------------------------------------
# TLS SNI extraction
# ---------------------------------------------------------------------------


def _maybe_extract_sni(
    tcp: Any,
    ts: float,
    sni_seen: dict[str, TlsSni],
    sample_cap: int,
) -> None:
    if sample_cap == 0:
        return  # 0 = disabled
    if len(sni_seen) >= sample_cap:
        return
    payload = bytes(tcp.data)
    if len(payload) < 5:
        return
    # TLS record: type(1) + version(2) + length(2). Type 22 = handshake.
    if payload[0] != 22:
        return
    # Inside the record, look for handshake type 1 = ClientHello.
    # Skip 5-byte record header → handshake header: type(1) + length(3).
    if len(payload) < 6:
        return
    if payload[5] != 1:  # not ClientHello
        return
    sni = _parse_sni_from_client_hello(payload[5:])
    if sni and sni not in sni_seen:
        sni_seen[sni] = TlsSni(
            hostname=sni,
            first_seen_at=datetime.fromtimestamp(ts, tz=UTC),
        )


def _parse_sni_from_client_hello(handshake: bytes) -> str | None:
    """Return the server_name extension's hostname, or None."""
    try:
        # Handshake: type(1) + length(3) + version(2) + random(32)
        if len(handshake) < 38:
            return None
        cursor = 1 + 3 + 2 + 32
        # session_id_length(1) + session_id
        sid_len = handshake[cursor]
        cursor += 1 + sid_len
        # cipher_suites_length(2) + cipher_suites
        if cursor + 2 > len(handshake):
            return None
        cs_len = struct.unpack("!H", handshake[cursor : cursor + 2])[0]
        cursor += 2 + cs_len
        # compression_methods_length(1) + compression_methods
        if cursor + 1 > len(handshake):
            return None
        cm_len = handshake[cursor]
        cursor += 1 + cm_len
        # extensions_length(2) + extensions
        if cursor + 2 > len(handshake):
            return None
        ext_len = struct.unpack("!H", handshake[cursor : cursor + 2])[0]
        cursor += 2
        ext_end = cursor + ext_len
        while cursor + 4 <= ext_end and cursor + 4 <= len(handshake):
            ext_type, ext_data_len = struct.unpack(
                "!HH",
                handshake[cursor : cursor + 4],
            )
            cursor += 4
            if ext_type == 0:  # server_name
                # server_name_list_length(2) + name_type(1) + name_length(2) + name
                if cursor + 5 > len(handshake):
                    return None
                # Skip list length + name type.
                name_len = struct.unpack(
                    "!H",
                    handshake[cursor + 3 : cursor + 5],
                )[0]
                name_start = cursor + 5
                if name_start + name_len > len(handshake):
                    return None
                return handshake[name_start : name_start + name_len].decode(
                    "ascii",
                    errors="replace",
                )
            cursor += ext_data_len
        return None
    except (IndexError, struct.error):
        return None


# ---------------------------------------------------------------------------
# HTTP extraction (plaintext request line only)
# ---------------------------------------------------------------------------


_HTTP_METHODS = frozenset({b"GET", b"POST", b"PUT", b"DELETE", b"HEAD", b"OPTIONS", b"PATCH"})


def _maybe_extract_http(
    tcp: Any,
    src: str,
    dst: str,
    ts: float,
    http_seen: list[HttpRequest],
    sample_cap: int,
) -> None:
    if sample_cap == 0:
        return  # 0 = disabled
    if len(http_seen) >= sample_cap:
        return
    payload = bytes(tcp.data)
    if not payload:
        return
    # Look for "METHOD path HTTP/" at the start of the payload.
    space = payload.find(b" ")
    if space < 0 or payload[:space] not in _HTTP_METHODS:
        return
    method = payload[:space].decode("ascii", errors="replace")
    rest = payload[space + 1 :]
    space2 = rest.find(b" ")
    if space2 < 0:
        return
    path = rest[:space2].decode("ascii", errors="replace")
    # Try to find a Host: header in the rest of the payload.
    host = ""
    host_idx = payload.lower().find(b"\r\nhost:")
    if host_idx >= 0:
        end = payload.find(b"\r\n", host_idx + 7)
        if end > 0:
            host = payload[host_idx + 7 : end].decode("ascii", errors="replace").strip()
    if not host:
        host = dst  # Fallback to destination IP.
    del src  # not surfaced in HttpRequest today; track via talkers
    http_seen.append(
        HttpRequest(
            method=method,
            host=host,
            path=path,
            seen_at=datetime.fromtimestamp(ts, tz=UTC),
        ),
    )


# ---------------------------------------------------------------------------
# pcap iteration (streaming)
# ---------------------------------------------------------------------------


def iter_packets(pcap_path: Path) -> Iterator[tuple[float, bytes]]:
    """Yield ``(timestamp, raw_bytes)`` for every packet in a pcap.

    Auto-detects pcap vs pcapng by file magic. Closes the file
    descriptor when iteration completes (or the consumer abandons
    the generator).
    """
    import dpkt

    with pcap_path.open("rb") as f:
        magic = f.read(4)
        f.seek(0)
        # pcap magic: 0xa1b2c3d4 / 0xd4c3b2a1 / 0xa1b23c4d / 0x4d3cb2a1
        # pcapng magic: 0x0a0d0d0a (block type SHB)
        if magic == b"\x0a\x0d\x0d\x0a":
            reader: Iterable[tuple[float, bytes]] = dpkt.pcapng.Reader(f)
        else:
            reader = dpkt.pcap.Reader(f)
        yield from reader


# ---------------------------------------------------------------------------
# IPv4/IPv6 helpers
# ---------------------------------------------------------------------------


def _ipv4_str(raw: bytes) -> str:
    import socket

    return socket.inet_ntop(socket.AF_INET, raw)


def _ipv6_str(raw: bytes) -> str:
    import socket

    return socket.inet_ntop(socket.AF_INET6, raw)
