#!/usr/bin/env python3
"""
netspeed - 实时网络速度监控工具 (htop 风格可视化)
从 /proc/net/dev 读取网络统计，无需 sudo
"""

import time
import sys
import os
import signal
import shutil
import argparse
import fnmatch

# ANSI 转义码
ESC = '\033'
RESET       = f'{ESC}[0m'
BOLD        = f'{ESC}[1m'
DIM         = f'{ESC}[2m'

# 标准色
RED         = f'{ESC}[31m'
GREEN       = f'{ESC}[32m'
YELLOW      = f'{ESC}[33m'
BLUE        = f'{ESC}[34m'
CYAN        = f'{ESC}[36m'
WHITE       = f'{ESC}[37m'

# 高亮色
HI_RED      = f'{ESC}[91m'
HI_GREEN    = f'{ESC}[92m'
HI_YELLOW   = f'{ESC}[93m'
HI_BLUE     = f'{ESC}[94m'
HI_CYAN     = f'{ESC}[96m'
HI_WHITE    = f'{ESC}[97m'

# 背景色
BG_GREEN    = f'{ESC}[42m'
BG_YELLOW   = f'{ESC}[43m'
BG_RED      = f'{ESC}[41m'
BG_BLUE     = f'{ESC}[44m'
BG_CYAN     = f'{ESC}[46m'

CLEAR       = f'{ESC}[2J{ESC}[H'
HIDE_CUR    = f'{ESC}[?25l'
SHOW_CUR    = f'{ESC}[?25h'

# 块字符
BLOCK_FULL  = '█'
BLOCK_HALF  = '▌'
BLOCK_EMPTY = '░'


def term_size():
    return shutil.get_terminal_size((80, 24))


def read_net_dev():
    """从 /proc/net/dev 读取网卡字节数"""
    stats = {}
    try:
        with open('/proc/net/dev', 'r') as f:
            for line in f.readlines()[2:]:
                parts = line.split(':')
                if len(parts) == 2:
                    iface = parts[0].strip()
                    vals = parts[1].split()
                    if len(vals) >= 9:
                        stats[iface] = {
                            'rx': int(vals[0]),
                            'tx': int(vals[8]),
                        }
    except OSError:
        pass
    return stats


def fmt_speed(bps):
    """格式化速度，固定宽度"""
    if bps >= 1024 ** 3:
        return f"{bps / 1024**3:7.2f} GB/s"
    elif bps >= 1024 ** 2:
        return f"{bps / 1024**2:7.2f} MB/s"
    elif bps >= 1024:
        return f"{bps / 1024:7.2f} KB/s"
    else:
        return f"{bps:7.0f}  B/s"


def fmt_total(b):
    """格式化总流量"""
    if b >= 1024 ** 4:
        return f"{b / 1024**4:.2f} TB"
    elif b >= 1024 ** 3:
        return f"{b / 1024**3:.2f} GB"
    elif b >= 1024 ** 2:
        return f"{b / 1024**2:.2f} MB"
    elif b >= 1024:
        return f"{b / 1024:.2f} KB"
    else:
        return f"{b} B"


def make_bar(value, peak, width):
    """
    绘制 htop 风格的彩色进度条：
    - 0%~50%  绿色
    - 50%~80% 黄色
    - 80%~100% 红色
    支持半块字符让分辨率翻倍
    """
    if peak <= 0:
        ratio = 0.0
    else:
        ratio = min(value / peak, 1.0)

    # 半块精度：每个字符占 0.5 格
    half_units = round(ratio * width * 2)
    full_blocks = half_units // 2
    has_half = half_units % 2

    bar = []
    for i in range(width):
        cell_ratio = (i + 0.5) / width
        # 根据位置决定颜色
        if cell_ratio < 0.5:
            color = HI_GREEN
        elif cell_ratio < 0.8:
            color = HI_YELLOW
        else:
            color = HI_RED

        if i < full_blocks:
            bar.append(f"{color}{BLOCK_FULL}{RESET}")
        elif i == full_blocks and has_half:
            bar.append(f"{color}{BLOCK_HALF}{RESET}")
        else:
            bar.append(f"{DIM}{BLOCK_EMPTY}{RESET}")

    return ''.join(bar)


class Monitor:
    def __init__(self, interval=1.0, filter_lo=True, filter_pattern=None):
        self.interval = interval
        self.filter_lo = filter_lo
        self.filter_pattern = filter_pattern  # None = 显示全部，否则按前缀或 glob 匹配
        self.prev = {}
        self.prev_time = None
        self.peak_rx = {}
        self.peak_tx = {}
        self.running = True

        signal.signal(signal.SIGINT, self._exit)
        signal.signal(signal.SIGTERM, self._exit)

    def _exit(self, *_):
        self.running = False
        sys.stdout.write(SHOW_CUR + '\n')
        sys.stdout.flush()
        sys.exit(0)

    def calc_speeds(self):
        """采样一次，返回各网卡速度 dict"""
        cur = read_net_dev()
        now = time.time()
        speeds = {}

        if self.prev and self.prev_time:
            dt = now - self.prev_time
            if dt <= 0:
                dt = 1e-9

            for iface, data in cur.items():
                if self.filter_lo and iface == 'lo':
                    continue
                if self.filter_pattern:
                    pat = self.filter_pattern
                    # 无通配符时做前缀匹配，有通配符时用 glob
                    if '*' in pat or '?' in pat or '[' in pat:
                        if not fnmatch.fnmatch(iface, pat):
                            continue
                    else:
                        if not iface.startswith(pat):
                            continue
                if iface not in self.prev:
                    continue

                rx_diff = max(0, data['rx'] - self.prev[iface]['rx'])
                tx_diff = max(0, data['tx'] - self.prev[iface]['tx'])

                rx_spd = rx_diff / dt
                tx_spd = tx_diff / dt

                self.peak_rx[iface] = max(self.peak_rx.get(iface, 0), rx_spd)
                self.peak_tx[iface] = max(self.peak_tx.get(iface, 0), tx_spd)

                speeds[iface] = {
                    'rx_spd':   rx_spd,
                    'tx_spd':   tx_spd,
                    'rx_total': data['rx'],
                    'tx_total': data['tx'],
                }

        self.prev = cur
        self.prev_time = now
        return speeds

    def render(self, speeds):
        cols, _ = term_size()
        out = [CLEAR]

        # ── 标题栏 ───────────────────────────────────────────
        ts = time.strftime("%Y-%m-%d %H:%M:%S")
        title = f"{BOLD}{HI_CYAN}NetSpeed{RESET}"
        hint  = f"{DIM}Ctrl+C 退出  │  间隔 {self.interval}s{RESET}"
        clock = f"{HI_WHITE}{BOLD}{ts}{RESET}"
        out.append(f" {title}  {hint}  {clock}")
        out.append(f"{DIM}{'─' * cols}{RESET}")
        out.append("")

        if not speeds:
            out.append(f"  {YELLOW}正在采集数据，请稍候…{RESET}")
        else:
            # 标签宽度 + 速度字符串宽度固定，其余给进度条
            label_w = 3          # "RX " / "TX "
            speed_w = 13         # " 1234.56 MB/s"
            bracket = 2          # "[ ]"
            bar_w = max(10, cols - label_w - speed_w - bracket - 4)

            for iface, d in speeds.items():
                rx  = d['rx_spd']
                tx  = d['tx_spd']

                # 峰值至少 100 KB/s，避免满格刷屏
                rx_pk = max(self.peak_rx.get(iface, 0), rx, 100 * 1024)
                tx_pk = max(self.peak_tx.get(iface, 0), tx, 100 * 1024)

                # ── 网卡名 ─────────────────────────────────
                out.append(f"  {BOLD}{HI_WHITE}{iface}{RESET}  "
                            f"{DIM}peak RX {fmt_speed(rx_pk).strip()}  "
                            f"TX {fmt_speed(tx_pk).strip()}{RESET}")

                # ── RX 条 ──────────────────────────────────
                rx_bar = make_bar(rx, rx_pk, bar_w)
                out.append(f"  {HI_GREEN}RX{RESET}[{rx_bar}] {HI_GREEN}{fmt_speed(rx)}{RESET}")

                # ── TX 条 ──────────────────────────────────
                tx_bar = make_bar(tx, tx_pk, bar_w)
                out.append(f"  {HI_YELLOW}TX{RESET}[{tx_bar}] {HI_YELLOW}{fmt_speed(tx)}{RESET}")

                # ── 累计流量 ───────────────────────────────
                out.append(f"     {DIM}累计  ↓{fmt_total(d['rx_total'])}  "
                            f"↑{fmt_total(d['tx_total'])}{RESET}")
                out.append("")

        # ── 图例 ────────────────────────────────────────────
        out.append(f"{DIM}{'─' * cols}{RESET}")
        out.append(f"  {DIM}图例: "
                   f"{HI_GREEN}██{RESET}{DIM} <50%  "
                   f"{HI_YELLOW}██{RESET}{DIM} 50~80%  "
                   f"{HI_RED}██{RESET}{DIM} >80%  "
                   f"(相对当前峰值){RESET}")

        sys.stdout.write('\n'.join(out))
        sys.stdout.flush()

    def run(self):
        sys.stdout.write(HIDE_CUR)
        sys.stdout.flush()
        try:
            while self.running:
                speeds = self.calc_speeds()
                self.render(speeds)
                time.sleep(self.interval)
        finally:
            sys.stdout.write(SHOW_CUR + '\n')
            sys.stdout.flush()


def main():
    parser = argparse.ArgumentParser(
        description='实时网络速度监控 (htop 风格)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  netspeed                   # 默认 1 秒刷新
  netspeed -i 0.5            # 0.5 秒刷新
  netspeed -f eno            # 只显示 eno 开头的网口
  netspeed -f 'veth*'        # glob 匹配（含通配符时需加引号）
  netspeed -i 2 --lo         # 包含 lo 接口
        """
    )
    parser.add_argument('-i', '--interval', type=float, default=1.0,
                        help='刷新间隔（秒），默认 1.0')
    parser.add_argument('-f', '--filter', dest='filter_pattern', default=None,
                        help='网口名称过滤：前缀匹配（如 eno）或 glob（如 veth*）')
    parser.add_argument('--lo', action='store_true',
                        help='包含回环接口 lo')
    args = parser.parse_args()

    Monitor(interval=args.interval, filter_lo=not args.lo,
            filter_pattern=args.filter_pattern).run()


if __name__ == '__main__':
    main()
