"""Per-entity stream state (tier 0): what one flow cannot show, the source's recent behaviour can.

For every source (and every destination) a small record of exponentially decayed counters is kept - flows,
bytes, packets; share of flows that got no reply, that were SYN-only, that were reset, that were very short;
destination fan-out (new port / new host within memory); repetition of the same (host, port); time since the
last flow. Every counter decays with half-life H, so "flows" reads as events in the last ~1.4 H seconds and all
shares are scale-free - they transfer between networks far better than absolute per-flow values.

Three horizons (default 10 s / 2 min / 30 min) see bursts, sessions and slow scans. Per flow the cost is O(1):
two dict lookups and ~60 multiply-adds. The state is bounded (LRU over `max_entities`), so a load-balancer
replica holds a few hundred MB at most. featurize() runs the same update over a time-sorted table (training), so
training features and runtime features are the same numbers.
"""
from __future__ import annotations

import math
from collections import OrderedDict

import numpy as np

from dibaguard.decide import evict

HORIZONS_S = (10.0, 120.0, 1800.0)
# authentication services: where repeated connections mean password guessing (also used by the campaign tier)
AUTH_PORTS = [21, 22, 23, 25, 110, 143, 389, 445, 465, 587, 993, 995, 1433, 1521, 3306, 3389, 5432, 5900, 6379, 9200, 27017]
# where a fast stream of repeated sessions *without* failure signs is still password guessing (directory and file services -
# LDAP, SMB - are chatty by nature: there, only failures make it brute force)
GUESSING_PORTS = [p for p in AUTH_PORTS if p not in (389, 445)]
PER_H = ["flows", "bytes_in", "pkts_in", "noreply", "synonly", "rst", "newport", "newhost", "repeat", "short"]
K = len(PER_H)
TAIL = ["gap_s", "gap_ewma_s", "dst_flows_h1", "dst_noreply_h1", "dst_newsrc_h1", "dst_flows_h2"]
MEMORY = 512          # (host, port) keys remembered per source; sources remembered per destination
# attack-direction features for the site-only anomaly tier: more of any of these is more suspicious (upper tail only);
# sizes and gaps are left out - benign bulk transfers and idle periods must not raise alarms; destination-side
# features are left out too - a surge of new clients on one host would flag every one of them (that signal belongs
# to the site pressure, not to a per-flow verdict)
ANOMALY = [f"src_{n}_h{i + 1}" for i in range(len(HORIZONS_S))
           for n in ("flows", "noreply", "synonly", "rst", "newport", "newhost", "repeat", "short")]
# the site-only anomaly tier leaves 'repeat' out: a benign client's repeat share climbs for hours as it revisits the
# same endpoints (it is an attack signal only together with failures, which the family rules check)
# the site-only anomaly tier judges the *shape* of a source's behaviour, not its volume: pure rate (src_flows_*) is
# left out - a busy legitimate client (an API integration, a heavy user, a busier hour) is not an attack by itself;
# volume attacks are the monotone prior's, the destination-flood tier's and the global tiers' job
ANOMALY_SITE = [f for f in ANOMALY if "repeat" not in f and "_flows_" not in f]
SHORT_MS = 10.0       # a flow shorter than this with <= 2 packets counts as 'short' (probe / failed connect)


def feature_names(horizons=HORIZONS_S) -> list[str]:
    out = []
    for i in range(len(horizons)):
        out += [f"src_{n}_h{i + 1}" for n in PER_H]
    return out + TAIL


class _Rec:
    __slots__ = ("t", "c", "seen", "hosts", "gap_ewma", "n")

    def __init__(self, t: float, n: int):
        self.t, self.c, self.seen, self.hosts, self.gap_ewma, self.n = t, [0.0] * n, OrderedDict(), {}, -1.0, 0


class EntityState:
    def __init__(self, horizons=HORIZONS_S, max_entities: int = 2_000_000, gap_alpha: float = 0.2):
        self.tau = [h / math.log(2) for h in horizons]
        self.nh, self.max, self.ga = len(horizons), max_entities, gap_alpha
        self.src: "OrderedDict[object, _Rec]" = OrderedDict()
        self.dst: "OrderedDict[object, _Rec]" = OrderedDict()

    def _get(self, table, key, t) -> tuple[_Rec, float, float]:
        """-> (record decayed to t, gap since its last event (-1 when new), lateness). A late record (older than the
        record's clock: a long flow exported at its end, another exporter's skewed clock) never rewinds the clock -
        that would make the next event decay everything accumulated - and is counted with the weight it has at the
        record's clock, exp(-lateness / tau): the exact update of a decayed counter for an out-of-order event"""
        r = table.get(key)
        if r is None:
            r = _Rec(t, self.nh * K)
            table[key] = r
            if len(table) > self.max:
                # a source with a history (8+ flows) is passed over while fresh one-flow sources can go instead
                evict(table, lambda v: v.n >= 8)
            return r, -1.0, 0.0
        table.move_to_end(key)
        dt = t - r.t
        if dt > 0:
            c = r.c
            for i, tau in enumerate(self.tau):
                f = math.exp(-dt / tau)
                for j in range(i * K, i * K + K):
                    c[j] *= f
            r.t = t
            return r, dt, 0.0
        return r, 0.0, -dt

    @staticmethod
    def _remember(rec: _Rec, key, host) -> bool:
        """Remember key; returns whether `host` was already known (before this flow)."""
        known = host in rec.hosts
        s = rec.seen
        if key in s:
            s.move_to_end(key)
        else:
            s[key] = host
            rec.hosts[host] = rec.hosts.get(host, 0) + 1
            if len(s) > MEMORY:
                _, h = s.popitem(last=False)
                n = rec.hosts[h] - 1
                if n:
                    rec.hosts[h] = n
                else:
                    del rec.hosts[h]
        return known

    def update(self, t: float, src, dst, dport: int, bytes_in: float, pkts_in: float, pkts_out: float,
               syn_only: bool, rst: bool, short: bool, out: list) -> None:
        """Append this flow's entity features (len(feature_names())) to `out`, then absorb the flow."""
        r, gap, late = self._get(self.src, src, t)
        d, _, dlate = self._get(self.dst, dst, t)
        r.n += 1                              # true age of the source (flows ever seen), not a decayed count
        key = (dst, dport)
        seen = key in r.seen
        new_host = dst not in r.hosts
        new_src = src not in d.seen
        noreply = pkts_out <= 0
        ev = (1.0, bytes_in, pkts_in, float(noreply), float(syn_only), float(rst), float(not seen), float(new_host),
              float(seen), float(short))
        c = r.c
        if late > 0 or dlate > 0:                 # out of order: each horizon counts it at its decayed weight
            self._update_late(r, d, ev, noreply, new_src, gap, late, dlate, out)
        else:
            for i in range(self.nh):
                b = i * K
                n = c[b] + 1.0
                out.append(n)
                out.append(c[b + 1] + bytes_in)
                out.append(c[b + 2] + pkts_in)
                for j in range(3, K):
                    out.append((c[b + j] + ev[j]) / n)
            out.append(gap)
            out.append(r.gap_ewma)
            dc = d.c
            n1 = dc[0] + 1.0
            out.append(n1)
            out.append((dc[3] + noreply) / n1)
            out.append((dc[6] + new_src) / n1)
            out.append(dc[K] + 1.0)
            for i in range(self.nh):
                b = i * K
                for j in range(K):
                    c[b + j] += ev[j]
                dc[b] += 1.0
                dc[b + 3] += noreply
                dc[b + 6] += new_src
        self._remember(r, key, dst)
        self._remember(d, src, src)
        if gap >= 0:
            r.gap_ewma = gap if r.gap_ewma < 0 else (1 - self.ga) * r.gap_ewma + self.ga * gap

    def _update_late(self, r: _Rec, d: _Rec, ev, noreply, new_src, gap, late, dlate, out: list) -> None:
        w = [max(math.exp(-late / tau), 1e-9) for tau in self.tau]
        dw = [max(math.exp(-dlate / tau), 1e-9) for tau in self.tau]
        c, dc = r.c, d.c
        for i in range(self.nh):
            b, wi = i * K, w[i]
            n = c[b] + wi
            out.append(n)
            out.append(c[b + 1] + wi * ev[1])
            out.append(c[b + 2] + wi * ev[2])
            for j in range(3, K):
                out.append((c[b + j] + wi * ev[j]) / n)
        out.append(gap)
        out.append(r.gap_ewma)
        n1 = dc[0] + dw[0]
        out.append(n1)
        out.append((dc[3] + dw[0] * noreply) / n1)
        out.append((dc[6] + dw[0] * new_src) / n1)
        out.append(dc[K] + dw[1])
        for i in range(self.nh):
            b = i * K
            for j in range(K):
                c[b + j] += w[i] * ev[j]
            dc[b] += dw[i]
            dc[b + 3] += dw[i] * noreply
            dc[b + 6] += dw[i] * new_src

    def featurize(self, t, src, dst, dport, bytes_in, pkts_in, pkts_out, syn_only, rst, duration_ms) -> np.ndarray:
        """Column arrays, already sorted by t -> float32 [n, len(feature_names())]. Runs the stream update."""
        n = len(t)
        short = (np.asarray(duration_ms, np.float64) < SHORT_MS) & (np.asarray(pkts_in, np.float64) <= 2)
        cols = [np.asarray(a).tolist() for a in (t, src, dst, dport, bytes_in, pkts_in, pkts_out, syn_only, rst, short)]
        out: list = []
        up = self.update
        for row in zip(*cols):
            up(*row, out)
        return np.asarray(out, np.float32).reshape(n, -1)


def tcp_syn_only(client_flags, server_flags) -> np.ndarray:
    """NetFlow TCP_FLAGS bitmask: client sent SYN (0x02) and no ACK (0x10), server nothing."""
    cf = np.asarray(client_flags, np.int64)
    return ((cf & 0x02) != 0) & ((cf & 0x10) == 0) & (np.asarray(server_flags, np.int64) == 0)


def tcp_rst(flags) -> np.ndarray:
    return (np.asarray(flags, np.int64) & 0x04) != 0


WELL_KNOWN = (20, 21, 22, 23, 25, 53, 67, 80, 110, 123, 135, 137, 139, 143, 161, 389, 443, 445, 465, 587, 993, 995, 1433, 1883,
              3306, 3389, 5432, 5900, 6379, 8080, 8443, 8883)
_WK = {p: i + 1 for i, p in enumerate(WELL_KNOWN)}


def service_key(dport, proto) -> np.ndarray:
    """Mondrian calibration group of a flow: protocol x destination service (well-known port exact, else 'other').
    A site's thresholds are kept per service, so a new or bursty service cannot move the others' thresholds."""
    dp = np.asarray(dport, np.int64)
    pr = np.asarray(proto, np.int64)
    svc = np.array([_WK.get(int(p), 0) if p >= 0 else 0 for p in dp], np.int64)
    svc = np.where((dp >= 1024), len(WELL_KNOWN) + 1, svc)
    return (np.clip(pr, 0, 255) * (len(WELL_KNOWN) + 2) + svc).astype(np.int32)


_F = {n: i for i, n in enumerate(feature_names())}


def family_rules(x, port=None) -> np.ndarray:
    """Transparent attack family from the entity stream features of alarmed flows ([n, 36] -> object array).
    Ranked by what the counters show over the 2-minute horizon: port walking, repeated failed connects to one
    service, floods at one target, host walking; 'attack_generic' when none of the signatures is clear."""
    x = np.asarray(x, np.float64)
    h = "h2"
    flows, nop, syn, newp, newh, rep, short = (x[:, _F[f"src_{k}_{h}"]] for k in ("flows", "noreply", "synonly", "newport", "newhost", "repeat", "short"))
    dflows, dnew = x[:, _F["dst_flows_h1"]], x[:, _F["dst_newsrc_h1"]]
    rate10 = x[:, _F["src_flows_h1"]]                       # flows in the last ~10 s
    out = np.full(len(x), "attack_generic", dtype=object)
    # the same (host, port) again and again - *and* failing (no reply / cut short / reset) or at an attacker's rate:
    # steady successful polling (telemetry, Kerberos / LDAP renewals, API clients, NTP) repeats too, and is not brute force
    rst = x[:, _F[f"src_rst_{h}"]]
    # (a fast repeater without failures is brute force only on an authentication service - many connections to one
    # proxy, API or broker are a busy client)
    auth = np.isin(np.asarray(port).astype(np.int64), GUESSING_PORTS) if port is not None else np.ones(len(x), bool)
    # ('short' - < 10 ms, <= 2 packets - is not a failure sign here: every DNS / NTP exchange over UDP is that short, and a
    # guessed password completes its TCP handshake; refused or unanswered connects show as reset / no reply)
    out[(rep > 0.6) & (flows >= 10) & ((nop > 0.2) | (rst > 0.2) | ((flows >= 30) & auth))] = "bruteforce"
    out[(newh > 0.3) & (flows >= 10) & (newp <= 0.5) & ((nop > 0.5) | (short > 0.5))] = "recon"
    # host walking on the services lateral movement uses (SMB, RPC, WinRM, SSH, RDP): one host reaching many others
    # where it has no business - answered or not - is a compromised host spreading, not a survey
    if port is not None:
        lat = np.isin(np.asarray(port).astype(np.int64), [445, 135, 139, 5985, 5986, 22, 3389])
        out[lat & (newh > 0.3) & (flows >= 10) & (newp <= 0.5)] = "lateral"
    # (a scan's probes go unanswered, half-open or refused; 'short' alone is every UDP DNS / NTP exchange - a resolver
    # asking many authoritative servers is not port walking)
    out[(newp > 0.5) & (flows >= 10) & ((syn > 0.5) | (nop > 0.5) | (rst > 0.5))] = "portscan"
    # a slow scan walks ports just the same over the 30-minute horizon (a probe every half minute never makes 10 flows
    # in 2 minutes): named when nothing else on the short horizon names it (playbook A10)
    newp3, flows3 = x[:, _F["src_newport_h3"]], x[:, _F["src_flows_h3"]]
    fail3 = (x[:, _F["src_synonly_h3"]] > 0.5) | (x[:, _F["src_noreply_h3"]] > 0.5) | (x[:, _F["src_short_h3"]] > 0.5)
    out[(out == "attack_generic") & (newp3 > 0.5) & (flows3 >= 10) & fail3] = "portscan"
    flood = (rate10 >= 100) & (newp <= 0.2) & (newh <= 0.2)  # >= 10 flows/s at one service for the whole horizon
    out[flood] = "dos"
    out[flood & (dnew > 0.5) & (dflows >= 500)] = "ddos"
    return out