"""Global download-bandwidth governor (one token bucket for the whole agent).

The agent lives on someone's desktop and shares that desktop's network —
typically WiFi, where saturating the link doesn't just slow other apps: it
bufferbloats the airtime (RTT 21→45ms) and collapses EVERY connection,
the operator's browser included. Measured 2026-07-02: the idle link moved
~38 MB/s single-stream, but under the 8-downloader sweep every stream on the
machine crawled at ~1-1.5 MB/s. So the cap is on the AGGREGATE, not per
stream: still downloads pump their chunks through take(), and ffmpeg video
streams — whose sockets live in a subprocess we can't wrap — are metered from
outside via /proc/<pid>/io and paused (SIGSTOP) into budget using charge()'s
debt signal; TCP flow control then stalls the sender while ffmpeg sleeps.

Accounting is post-paid (charge the bytes first, then wait out any debt): the
bytes have already crossed the network by the time we count them, and it means
a chunk larger than one second of budget can never deadlock the bucket.
Stdlib-only on purpose — unit-tested in CI, where the agent's ML deps
don't exist.
"""
import threading
import time


class TokenBucket:
    """Thread-safe token bucket in bytes/second. rate 0 = unlimited.

    `consumed` is the monotonic total of bytes charged (throttled or not) —
    the worker's rate loop derives the UI's "net MB/s" readout from it.
    """

    def __init__(self, rate_bytes_per_s: float = 0.0):
        self._cond = threading.Condition()
        self._rate = max(0.0, float(rate_bytes_per_s))
        # Burst = one second of budget: enough that chunked reads stay smooth,
        # small enough that a burst can't meaningfully lift the average.
        self._level = self._rate
        self._stamp = time.monotonic()
        self.consumed = 0

    @property
    def rate(self) -> float:
        return self._rate

    def set_rate(self, rate_bytes_per_s: float) -> None:
        """Retune live (the UI dial). Waiters re-check immediately, so raising
        the cap (or lifting it with 0) unblocks a mid-download wait at once."""
        with self._cond:
            self._refill_locked()          # settle elapsed time at the OLD rate
            self._rate = max(0.0, float(rate_bytes_per_s))
            self._level = min(self._level, self._rate)
            self._cond.notify_all()

    def _refill_locked(self) -> None:
        now = time.monotonic()
        self._level = min(self._rate, self._level + (now - self._stamp) * self._rate)
        self._stamp = now

    def take(self, n: int) -> None:
        """Charge n bytes and block until the budget recovers (stills path)."""
        with self._cond:
            self.consumed += n
            if self._rate <= 0:
                return
            self._refill_locked()
            self._level -= n
            while self._level < 0:
                # Wake early on set_rate; cap the wait so a big debt is paid in
                # re-checked slices rather than one long uninterruptible sleep.
                self._cond.wait(min(-self._level / self._rate, 0.5))
                if self._rate <= 0:
                    return
                self._refill_locked()

    def charge(self, n: int) -> float:
        """Charge n bytes WITHOUT blocking; return seconds of debt (0 = within
        budget). The ffmpeg governor can't block the subprocess's own reads, so
        it SIGSTOPs the process for (about) the returned debt instead."""
        with self._cond:
            self.consumed += n
            if self._rate <= 0:
                return 0.0
            self._refill_locked()
            self._level -= n
            return max(0.0, -self._level / self._rate)


class PidReadMeter:
    """Cumulative read-bytes meter for a subprocess, via /proc/<pid>/io.

    `rchar` counts every read() syscall's bytes — for a streaming ffmpeg the
    network reads dominate, so the delta is a good-enough aggregate-bandwidth
    signal (it's a governor, not a billing meter). Returns None when /proc is
    unavailable (process exited, or a non-Linux host): the caller then simply
    doesn't govern — degrade to unthrottled rather than break video sampling.
    """

    def __init__(self, pid: int):
        self._path = f"/proc/{pid}/io"
        self._last = 0

    def delta(self) -> int | None:
        try:
            with open(self._path, "rb") as f:
                for line in f:
                    if line.startswith(b"rchar:"):
                        total = int(line.split()[1])
                        d, self._last = total - self._last, total
                        return max(0, d)
        except (OSError, ValueError):
            return None
        return None