"""Flow feature views and raw label -> family mapping, shared by the corpus builder, the trainer and the runtime.

A view is a fixed, ordered list of base columns plus a few derived scale-free ratios (computed here, so training and
inference can never disagree). Identity columns are never features: addresses, OS-chosen source ports, timestamps,
flow / DNS query ids describe *who*, not *how* - the classic shortcut that makes flow models score 99 % inside one
dataset and collapse on a new network. SHORTCUTS lists the host-fingerprint columns (TTL = hop distance of the attack
machine, initial TCP window = its OS) that the 'robust' model variant leaves out; the trainer measures both.
"""
from __future__ import annotations

import numpy as np

NF = ["L4_DST_PORT", "PROTOCOL", "L7_PROTO", "IN_BYTES", "IN_PKTS", "OUT_BYTES", "OUT_PKTS", "TCP_FLAGS",
      "CLIENT_TCP_FLAGS", "SERVER_TCP_FLAGS", "FLOW_DURATION_MILLISECONDS", "DURATION_IN", "DURATION_OUT", "MIN_TTL",
      "MAX_TTL", "LONGEST_FLOW_PKT", "SHORTEST_FLOW_PKT", "MIN_IP_PKT_LEN", "MAX_IP_PKT_LEN",
      "SRC_TO_DST_SECOND_BYTES", "DST_TO_SRC_SECOND_BYTES", "RETRANSMITTED_IN_BYTES", "RETRANSMITTED_IN_PKTS",
      "RETRANSMITTED_OUT_BYTES", "RETRANSMITTED_OUT_PKTS", "SRC_TO_DST_AVG_THROUGHPUT", "DST_TO_SRC_AVG_THROUGHPUT",
      "NUM_PKTS_UP_TO_128_BYTES", "NUM_PKTS_128_TO_256_BYTES", "NUM_PKTS_256_TO_512_BYTES",
      "NUM_PKTS_512_TO_1024_BYTES", "NUM_PKTS_1024_TO_1514_BYTES", "TCP_WIN_MAX_IN", "TCP_WIN_MAX_OUT", "ICMP_TYPE",
      "DNS_QUERY_TYPE", "DNS_TTL_ANSWER", "FTP_COMMAND_RET_CODE"]

# CICFlowMeter columns common to the DistriNet-corrected CIC-IDS2017 / CSE-CIC-IDS2018 and CIC-DDoS2019 releases
CIC = ["Protocol", "Flow Duration", "Total Fwd Packets", "Total Backward Packets", "Fwd Packets Length Total",
       "Bwd Packets Length Total", "Fwd Packet Length Max", "Fwd Packet Length Min", "Fwd Packet Length Mean",
       "Fwd Packet Length Std", "Bwd Packet Length Max", "Bwd Packet Length Min", "Bwd Packet Length Mean",
       "Bwd Packet Length Std", "Flow Bytes/s", "Flow Packets/s", "Flow IAT Mean", "Flow IAT Std", "Flow IAT Max",
       "Flow IAT Min", "Fwd IAT Total", "Fwd IAT Mean", "Fwd IAT Std", "Fwd IAT Max", "Fwd IAT Min", "Bwd IAT Total",
       "Bwd IAT Mean", "Bwd IAT Std", "Bwd IAT Max", "Bwd IAT Min", "Fwd PSH Flags", "Bwd PSH Flags", "Fwd URG Flags",
       "Bwd URG Flags", "Fwd Header Length", "Bwd Header Length", "Fwd Packets/s", "Bwd Packets/s",
       "Packet Length Min", "Packet Length Max", "Packet Length Mean", "Packet Length Std", "Packet Length Variance",
       "FIN Flag Count", "SYN Flag Count", "RST Flag Count", "PSH Flag Count", "ACK Flag Count", "URG Flag Count",
       "ECE Flag Count", "Down/Up Ratio", "Avg Packet Size", "Avg Fwd Segment Size", "Avg Bwd Segment Size",
       "Fwd Avg Bytes/Bulk", "Fwd Avg Packets/Bulk", "Fwd Avg Bulk Rate", "Bwd Avg Bytes/Bulk",
       "Bwd Avg Packets/Bulk", "Bwd Avg Bulk Rate", "Subflow Fwd Packets", "Subflow Fwd Bytes",
       "Subflow Bwd Packets", "Subflow Bwd Bytes", "Init Fwd Win Bytes", "Init Bwd Win Bytes",
       "Fwd Act Data Packets", "Fwd Seg Size Min", "Active Mean", "Active Std", "Active Max", "Active Min",
       "Idle Mean", "Idle Std", "Idle Max", "Idle Min"]

LU = ["avg_ipt", "bytes_in", "bytes_out", "entropy", "num_pkts_out", "num_pkts_in", "proto", "total_entropy",
      "duration"]

BASE = {"nf": NF, "cic": CIC, "lu": LU}
# derived = (name, numerator, denominator): scale-free ratios trees cannot form from single-column splits
DERIVED = {
    "nf": [("in_bytes_per_pkt", "IN_BYTES", "IN_PKTS"), ("out_bytes_per_pkt", "OUT_BYTES", "OUT_PKTS"),
           ("pkt_ratio", "OUT_PKTS", "IN_PKTS"), ("byte_ratio", "OUT_BYTES", "IN_BYTES"),
           ("ms_per_pkt", "FLOW_DURATION_MILLISECONDS", "IN_PKTS")],
    "cic": [("byte_ratio", "Bwd Packets Length Total", "Fwd Packets Length Total"),
            ("hdr_per_fwd_pkt", "Fwd Header Length", "Total Fwd Packets")],
    "lu": [("in_bytes_per_pkt", "bytes_in", "num_pkts_in"), ("out_bytes_per_pkt", "bytes_out", "num_pkts_out"),
           ("pkt_ratio", "num_pkts_out", "num_pkts_in"), ("byte_ratio", "bytes_out", "bytes_in")],
}
SHORTCUTS = {"nf": ["MIN_TTL", "MAX_TTL", "TCP_WIN_MAX_IN", "TCP_WIN_MAX_OUT"],
             "cic": ["Init Fwd Win Bytes", "Init Bwd Win Bytes", "Fwd Seg Size Min"],
             "lu": []}
BIG = np.float32(1e30)


def feature_names(view: str, robust: bool = False) -> list[str]:
    drop = set(SHORTCUTS[view]) if robust else set()
    return [c for c in BASE[view] if c not in drop] + [d[0] for d in DERIVED[view]]


def matrix(view: str, cols, robust: bool = False) -> np.ndarray:
    """cols: mapping column name -> 1-d array (a DataFrame works) -> float32 [n, len(feature_names(view, robust))].
    Missing values -> -1, infinities clipped (CICFlowMeter writes inf for zero-duration rates)."""
    drop = set(SHORTCUTS[view]) if robust else set()
    out = [np.asarray(cols[c], np.float64) for c in BASE[view] if c not in drop]
    for _, num, den in DERIVED[view]:
        a, b = np.asarray(cols[num], np.float64), np.asarray(cols[den], np.float64)
        out.append(a / np.maximum(b, 1.0))
    x = np.column_stack(out) if out else np.zeros((0, 0))
    x = np.nan_to_num(x, nan=-1.0, posinf=float(BIG), neginf=-float(BIG))
    return np.clip(x, -BIG, BIG).astype(np.float32)


# raw label (lower-cased) -> family in families.SEVERITY; None = not used (neither clean benign nor a real attack)
_FAMILY = {
    "benign": "benign", "normal": "benign", "normal flow": "benign",
    "syn scan - aggressive": "portscan", "syn scan": "portscan", "xmas scan": "portscan", "fin scan": "portscan",
    "null scan": "portscan", "udp scan": "portscan", "tcp connect scan": "portscan", "ack scan": "portscan",
    "attempted-relabel-as-benign": None,   # DistriNet: an attack attempt that did not execute - excluded both ways
    "outlier": None,                       # LUFlow: unexplained, not labelled either way
    "malicious": "attack_generic",         # LUFlow: matched to honeypot / threat-intel sources
    "ddos": "ddos", "dos": "dos", "heartbleed": "exploit",
    "scanning": "portscan", "portscan": "portscan", "infiltration - portscan": "portscan",
    "infiltration - nmap portscan": "portscan", "reconnaissance": "recon", "analysis": "recon",
    "password": "bruteforce", "brute force": "bruteforce", "ftp-patator": "bruteforce", "ssh-patator": "bruteforce",
    "ssh-bruteforce": "bruteforce", "web attack - brute force": "bruteforce",
    "xss": "webattack", "injection": "webattack", "web attack - xss": "webattack",
    "web attack - sql injection": "webattack", "web attack - sql": "webattack",
    "bot": "botnet", "botnet": "botnet", "botnet ares": "botnet",
    "infilteration": "infiltration", "infiltration": "infiltration",
    "infiltration - dropbox download": "infiltration", "infiltration - communication victim attacker": "infiltration",
    "exploits": "exploit", "shellcode": "exploit", "backdoor": "backdoor", "fuzzers": "fuzzing",
    "generic": "attack_generic", "worms": "worm", "ransomware": "malware", "theft": "exfil", "mitm": "mitm",
}
_PREFIX = [("dos ", "dos"), ("dos-", "dos"), ("dos_", "dos"), ("ddos", "ddos"), ("drdos_", "ddos"), ("recon", "recon"),
           ("port scan", "portscan"), ("portscan", "portscan"), ("ping sweep", "recon"), ("os scan", "recon"),
           ("vulnerability scan", "recon"), ("arp spoof", "mitm"), ("spoof", "mitm"), ("mitm", "mitm"),
           ("mqtt", "dos"), ("brute", "bruteforce"), ("dictionary", "bruteforce"), ("slowloris", "dos"),
           ("mirai", "botnet"), ("benign", "benign")]
# CIC-DDoS2019 attack names (all reflection / flooding DDoS)
_DDOS2019 = {"syn", "udp", "udp-lag", "udplag", "tftp", "ldap", "mssql", "netbios", "portmap", "webddos", "snmp",
             "dns", "ntp", "ssdp"}


def family(raw: str, testbed: str = "") -> str | None:
    r = str(raw).strip().lower()
    if r in _FAMILY:
        return _FAMILY[r]
    for p, f in _PREFIX:
        if r.startswith(p):
            return f
    if testbed == "ddos2019" and r in _DDOS2019:
        return "ddos"
    return "attack_generic"