#!/usr/bin/env python3
"""
Shared rate limiter for NetVM scripts.

Any .py script can use this to avoid hitting Muse rate limits:
    from rate_limiter import rate_limit_wait
    rate_limit_wait("opm")  # blocks until allowed

Uses a token bucket per agent stored in /tmp (persists across invocations,
not across reboots). Default: 1 op per 3s sustained, burst of 5, max ~20/min.

Jitter (2026-10-04): each agent gets a DETERMINISTIC per-agent jitter factor
derived from sha256(salt + agent), in [1-jitter, 1+jitter] (default
jitter=0.3, i.e. +/-30%). Rationale: all 4 fleet nodes egress from a single
Cloudflare IP, so Cloudflare sees correlated traffic. When fleet timers fire
simultaneously, per-agent jitter drifts each node's send phase apart instead
of hitting in lockstep. Deterministic per agent = reproducible behavior;
different agents = de-correlated phases. Does not make limits stricter: the
mean interval is unchanged.
"""
import hashlib
import json
import os
import random
import time

RATE_LIMIT_FILE = "/tmp/netvm-rate-limit.json"
RATE_LIMIT_INTERVAL = 3.0  # seconds between ops (base; jittered per agent)
RATE_LIMIT_BURST = 5
RATE_LIMIT_MAX_PER_MIN = 20
DEFAULT_JITTER = 0.3  # +/-30% deterministic per-agent interval jitter

JITTER_SEED_SALT = "netvm-rate-jitter:v1:"


def _jitter_factor(agent, jitter=DEFAULT_JITTER):
    """Deterministic per-agent multiplier in [1-jitter, 1+jitter].

    Seeded by sha256(salt + agent): the same agent always gets the same
    factor (reproducible), different agents get different factors
    (de-correlated). Uses an isolated Random instance; global random
    state is untouched.
    """
    if jitter <= 0:
        return 1.0
    digest = hashlib.sha256((JITTER_SEED_SALT + agent).encode()).digest()
    seed = int.from_bytes(digest[:8], "big")
    rng = random.Random(seed)
    return 1.0 + jitter * (rng.random() * 2.0 - 1.0)


def effective_interval(agent, interval=RATE_LIMIT_INTERVAL,
                       jitter=DEFAULT_JITTER):
    """The actual minimum op spacing for this agent after jitter."""
    return interval * _jitter_factor(agent, jitter)


def rate_limit_policy(agent, interval=RATE_LIMIT_INTERVAL,
                      burst=RATE_LIMIT_BURST, jitter=DEFAULT_JITTER):
    """Return the effective rate-limit policy for an agent (audit/docs)."""
    return {
        "agent": agent,
        "base_interval_s": interval,
        "jitter": jitter,
        "jitter_factor": round(_jitter_factor(agent, jitter), 4),
        "effective_interval_s": round(
            effective_interval(agent, interval, jitter), 3),
        "burst": burst,
        "max_per_min": RATE_LIMIT_MAX_PER_MIN,
        "state_file": RATE_LIMIT_FILE,
        "scope": ("per-agent buckets; state shared in one file, "
                  "keys namespaced by agent"),
    }


def _load_state():
    try:
        with open(RATE_LIMIT_FILE) as f:
            return json.load(f)
    except:
        return {}


def _save_state(state):
    try:
        # Atomic write via temp file
        tmp = RATE_LIMIT_FILE + ".tmp"
        with open(tmp, "w") as f:
            json.dump(state, f)
        os.rename(tmp, RATE_LIMIT_FILE)
    except:
        pass


def rate_limit_wait(agent, interval=RATE_LIMIT_INTERVAL,
                    burst=RATE_LIMIT_BURST, jitter=DEFAULT_JITTER):
    """
    Block until the agent is allowed to perform an operation.
    Returns the time waited in seconds (0 if no wait needed).

    The interval is jittered deterministically per agent so fleet nodes
    don't send in lockstep.
    """
    eff_interval = effective_interval(agent, interval, jitter)
    state = _load_state()
    now = time.time()

    # Prune old entries
    recent_key = f"{agent}_recent"
    recent = state.get(recent_key, [])
    recent = [t for t in recent if now - t < 60]

    # Check burst limit (max per minute)
    if len(recent) >= RATE_LIMIT_MAX_PER_MIN:
        wait = 60 - (now - recent[0]) + 1
        if wait > 0:
            time.sleep(wait)
            # Refresh after wait
            state = _load_state()
            recent = state.get(recent_key, [])
            recent = [t for t in recent if time.time() - t < 60]

    # Check interval limit (jittered per agent)
    last = state.get(agent, 0)
    now = time.time()
    if now - last < eff_interval:
        wait = eff_interval - (now - last)
        time.sleep(wait)

    # Record this operation
    now = time.time()
    state[agent] = now
    recent.append(now)
    state[recent_key] = recent
    _save_state(state)
    return 0


def rate_limit_check(agent, interval=RATE_LIMIT_INTERVAL,
                     jitter=DEFAULT_JITTER):
    """
    Non-blocking check. Returns (allowed: bool, wait_seconds: float).
    """
    eff_interval = effective_interval(agent, interval, jitter)
    state = _load_state()
    now = time.time()
    recent = state.get(f"{agent}_recent", [])
    recent = [t for t in recent if now - t < 60]

    if len(recent) >= RATE_LIMIT_MAX_PER_MIN:
        wait = 60 - (now - recent[0]) + 1
        return False, max(0, wait)

    last = state.get(agent, 0)
    if now - last < eff_interval:
        return False, eff_interval - (now - last)

    return True, 0


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser(
        description="Inspect NetVM per-agent rate-limit policy")
    p.add_argument("--policy", nargs="*", default=["muse", "pip", "646", "opm"],
                   help="Show effective policy per agent")
    p.add_argument("--no-jitter", action="store_true",
                   help="Show policy without jitter")
    args = p.parse_args()
    j = 0.0 if args.no_jitter else DEFAULT_JITTER
    for a in args.policy:
        print(json.dumps(rate_limit_policy(a, jitter=j), indent=2))