"""线程安全的本地转发流量统计。"""

import threading
import time


class TrafficStats:
    """记录单次代理会话的上下行总量与采样速率。"""

    def __init__(self, clock=time.time):
        self._clock = clock
        self._lock = threading.Lock()
        self._up = 0
        self._down = 0
        self._prev_up = 0
        self._prev_down = 0
        self._prev_ts = 0.0

    def add(self, up=0, down=0):
        with self._lock:
            self._up += up
            self._down += down

    def reset(self):
        with self._lock:
            self._up = self._down = 0
            self._prev_up = self._prev_down = 0
            self._prev_ts = self._clock()

    def rate(self):
        with self._lock:
            now = self._clock()
            if self._prev_ts > 0 and now > self._prev_ts:
                elapsed = now - self._prev_ts
                up_rate = int((self._up - self._prev_up) / elapsed)
                down_rate = int((self._down - self._prev_down) / elapsed)
            else:
                up_rate = down_rate = 0
            self._prev_up = self._up
            self._prev_down = self._down
            self._prev_ts = now
            return up_rate, down_rate, self._up, self._down

    def snapshot(self):
        with self._lock:
            return self._up, self._down


def format_bytes(value):
    """将字节数格式化为 UI 展示文本。"""
    number = float(value)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if number < 1024 or unit == "TB":
            return f"{int(number)} {unit}" if unit == "B" else f"{number:.1f} {unit}"
        number /= 1024