"""Per-NIC connection dynamics accumulator.

Passive model of path RTT, jitter and loss observed across network
operations on a NIC.  Operations report observations; new operations
query for recommended timeouts and probe counts.

Phase 1 (this file) uses a two-bucket pathway taxonomy:

  PATHWAY_LAN  - loopback, RFC1918 v4, link-local v4/v6, ULA v6,
                 multicast (treated as local-network for purposes
                 of timing assumptions).
  PATHWAY_WAN  - everything else.

Bucket key: ``(pathway, af, proto)``.  Phase-2 work could split WAN
further by destination AS / region, but two-buckets already captures
the LAN-vs-WAN constant-split that the codebase has been quietly
re-implementing in many places (FAST_PUNCH_PARAMS vs UDP_PUNCH_PARAMS,
TCP_PUNCH_LAN vs TCP_PUNCH_REMOTE, etc.).

RTT smoothing follows RFC 6298:
    SRTT'   = (1 - ALPHA) * SRTT + ALPHA * sample
    RTTVAR' = (1 - BETA)  * RTTVAR + BETA  * |SRTT - sample|
    RTO     = SRTT + 4 * RTTVAR
The recommended timeout multiplies RTO by ``RTO_SAFETY_FACTOR`` and
clamps to [MIN_TIMEOUT_S, MAX_TIMEOUT_S].

Loss rate is an EWMA of "was the operation lost / timed out" booleans.

Cold-start: a bucket with zero samples returns the caller-supplied
``default`` for both ``timeout_s`` and ``probe_count`` so the rest of
the system observes no behavioural change until enough data has
accumulated to make adaptive recommendations.

This module is intentionally PASSIVE: nothing in it issues network ops
on its own.  Callers thread observations in from their existing probes
(see e.g. ``nat_utils.delta_test`` for the first integration).
"""
from __future__ import print_function

import ipaddress


PATHWAY_LAN = "lan"
PATHWAY_WAN = "wan"

# RFC 6298 RTT smoothing weights.
ALPHA = 0.125
BETA = 0.25

# EWMA weight for loss-rate updates.
LOSS_ALPHA = 0.2

# Cold-start defaults (returned when bucket has no samples).
DEFAULT_TIMEOUT_S = 2.0
DEFAULT_PROBE_COUNT = 8

# Hard clamps on returned recommendations.
MIN_TIMEOUT_S = 0.5
MAX_TIMEOUT_S = 5.0
MIN_PROBE_COUNT = 6
MAX_PROBE_COUNT = 24

# Safety multiplier applied to (SRTT + 4*RTTVAR) before clamping.
# Above 1.0 makes the recommendation more forgiving than a strict
# RFC 6298 RTO; we're sizing probe budgets, not retransmission timers.
RTO_SAFETY_FACTOR = 1.5

# Number of samples below which we still partially trust the default.
# Prevents a single noisy first sample from dominating.
MIN_SAMPLES_FOR_FULL_TRUST = 4


def classify_pathway(dest_ip):
    """Bucket an IP address (string) into PATHWAY_LAN or PATHWAY_WAN.

    Unparseable inputs default to WAN (conservative timeouts).
    """
    if dest_ip is None:
        return PATHWAY_WAN
    try:
        addr = ipaddress.ip_address(_to_text(dest_ip))
    except (ValueError, TypeError):
        return PATHWAY_WAN
    if addr.is_loopback:
        return PATHWAY_LAN
    if addr.is_link_local:
        return PATHWAY_LAN
    if addr.is_private:
        return PATHWAY_LAN
    if addr.is_multicast:
        return PATHWAY_LAN
    return PATHWAY_WAN


def _to_text(value):
    """Coerce bytes or str to text without using f-strings (py35 compat)."""
    if isinstance(value, bytes):
        return value.decode("ascii", errors="ignore")
    return value


class Bucket:
    """One (pathway, af, proto) bucket's running statistics."""

    def __init__(self):
        self.samples = 0
        self.srtt_s = None    # smoothed RTT in seconds
        self.rttvar_s = None  # RTT variance estimate in seconds
        self.loss_rate = 0.0  # EWMA of "lost" booleans, 0..1

    def observe_ok(self, rtt_s):
        """Record a successful operation with measured RTT (seconds)."""
        if rtt_s < 0:
            return
        if self.srtt_s is None:
            self.srtt_s = rtt_s
            self.rttvar_s = rtt_s / 2.0
        else:
            err = abs(self.srtt_s - rtt_s)
            self.rttvar_s = (1.0 - BETA) * self.rttvar_s + BETA * err
            self.srtt_s = (1.0 - ALPHA) * self.srtt_s + ALPHA * rtt_s
        self.samples += 1
        # Decay loss rate on a successful sample.
        self.loss_rate = (1.0 - LOSS_ALPHA) * self.loss_rate

    def observe_loss(self):
        """Record a lost / timed-out operation."""
        self.samples += 1
        self.loss_rate = LOSS_ALPHA + (1.0 - LOSS_ALPHA) * self.loss_rate

    def rto_s(self):
        """RFC 6298 RTO in seconds, or None if no RTT data yet."""
        if self.srtt_s is None:
            return None
        return self.srtt_s + 4.0 * self.rttvar_s

    def as_dict(self):
        return {
            "samples": self.samples,
            "srtt_ms": int(self.srtt_s * 1000) if self.srtt_s is not None else None,
            "rttvar_ms": int(self.rttvar_s * 1000) if self.rttvar_s is not None else None,
            "loss_rate": round(self.loss_rate, 3),
        }


class ConnectionDynamics:
    """Accumulator across all (pathway, af, proto) buckets for one NIC."""

    def __init__(self, nic_name=None):
        self.nic_name = nic_name
        self.buckets = {}

    def bucket_for(self, dest_ip, af, proto):
        pathway = classify_pathway(dest_ip)
        key = (pathway, af, proto)
        b = self.buckets.get(key)
        if b is None:
            b = Bucket()
            self.buckets[key] = b
        return b

    def observe_ok(self, dest_ip, af, proto, rtt_s):
        self.bucket_for(dest_ip, af, proto).observe_ok(rtt_s)

    def observe_loss(self, dest_ip, af, proto):
        self.bucket_for(dest_ip, af, proto).observe_loss()

    def timeout_s(self, dest_ip, af, proto, default=DEFAULT_TIMEOUT_S):
        """Recommended timeout in seconds for an op on this pathway.

        Cold-start (zero samples): returns ``default``.  Once samples
        accumulate, returns clamp(SRTT + 4*RTTVAR, MIN, MAX) * safety.
        """
        b = self.bucket_for(dest_ip, af, proto)
        rto = b.rto_s()
        if rto is None or b.samples < MIN_SAMPLES_FOR_FULL_TRUST:
            # Blend in the default while building confidence.
            if rto is None:
                return default
            blend = b.samples / float(MIN_SAMPLES_FOR_FULL_TRUST)
            adaptive = rto * RTO_SAFETY_FACTOR
            value = blend * adaptive + (1.0 - blend) * default
        else:
            value = rto * RTO_SAFETY_FACTOR
        return max(MIN_TIMEOUT_S, min(MAX_TIMEOUT_S, value))

    def probe_count(self, dest_ip, af, proto, default=DEFAULT_PROBE_COUNT):
        """Recommended probe count for the pathway.

        Boosts when observed loss rate is high so the expected yield
        after losses still approximates ``default``.  Cold-start
        returns ``default`` unchanged.
        """
        b = self.bucket_for(dest_ip, af, proto)
        if b.samples == 0:
            return default
        if b.loss_rate <= 0.05:
            return default
        survival = max(0.3, 1.0 - b.loss_rate)
        scaled = int(round(default / survival))
        return max(MIN_PROBE_COUNT, min(MAX_PROBE_COUNT, scaled))

    def stats(self):
        """Return a dict of bucket-key -> stats dict for debug logging."""
        out = {}
        for key, b in self.buckets.items():
            out[key] = b.as_dict()
        return out


# Module-level registry so repeated NIC loads in the same process keep
# their learnings.  In production a NIC object is constructed once per
# node lifetime so this is mostly a convenience for the experimental
# scripts (which build Interface fresh each iteration); production code
# can also keep its own reference via Interface.connection_dynamics.
GLOBAL_REGISTRY = {}


def for_nic(nic_name):
    """Return a (persistent across re-creations) ConnectionDynamics for nic_name."""
    if nic_name is None:
        nic_name = "(default)"
    d = GLOBAL_REGISTRY.get(nic_name)
    if d is None:
        d = ConnectionDynamics(nic_name=nic_name)
        GLOBAL_REGISTRY[nic_name] = d
    return d


