#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Royirin Rosmontis Private Proxy JP 流量监控（Azure 门户口径版）
- 主计量口径：Azure 门户免费服务页 "Data Transfer Out (GB)"（Standard Data Transfer Out - Free，MGN 路由）
  —— 通过 VM 托管身份调用 Consumption Usage Details API 拉取，按自然月聚合，缓存至 data/azure_usage.json
- 参考口径：vnstat -i eth0（实时原始出站，仅作参考展示）
- 数据源（参考）：vnstat -i <interface> --json m/d（eth0 为 bond 主接口，覆盖 nginx 443 与 sing-box 8443 全部流量）
- 按 年/月 聚合成 JSON 存入 data/YYYY/YYYY-MM.json（弃用 SQLite）
- HTTP 服务: GET /vpn/status/traffic 返回当月用量 JSON（Azure 口径为主）
- 阈值告警: 达到配额 50/80/90/95/99/100% 时发邮件（每档每月只发一次，跳档补发；基于 Azure 口径）
- 月报: 每月 1 日 00:05 起发送上个月报告（失败重试；基于 Azure 口径）
"""
import argparse
import json
import logging
import os
import socket
import smtplib
import ssl
import subprocess
import sys
import threading
import time
import urllib.parse
import urllib.request
from datetime import date, datetime, timedelta, timezone
from email.message import EmailMessage
from email.utils import formataddr
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

log = logging.getLogger('traffic_monitor')

GB = 1_000_000_000  # Azure 数据转移按十进制 GB 计费


def utcnow():
    return datetime.now(timezone.utc)


def load_json(path, default):
    try:
        with open(path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        return default
    except (json.JSONDecodeError, OSError) as e:
        log.warning('read %s failed: %s; using default', path, e)
        return default


def save_json(path, obj):
    tmp = path + '.tmp'
    with open(tmp, 'w', encoding='utf-8') as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)
    os.replace(tmp, path)


def gb_str(n_bytes):
    return f'{n_bytes / GB:.2f} GB'


def month_struct(year, month):
    return {
        'year': int(year),
        'month': int(month),
        'rx_bytes': 0,
        'tx_bytes': 0,
        'days': [],
        'updated_at': utcnow().isoformat(),
    }


class Monitor:
    def __init__(self, config):
        self.config = config
        self.lock = threading.RLock()
        self.months = {}
        self.alert_state = None

        os.makedirs(self.config['data_dir'], exist_ok=True)
        self._load_all_months()
        self._load_alert_state()
        # Azure 计费用量（Usage Details API，VM 托管身份鉴权）
        self.azure_usage = None
        self._azure_last_fetch = 0.0
        cache_path = self.config.get('azure_usage_cache_path')
        if cache_path and os.path.exists(cache_path):
            self.azure_usage = load_json(cache_path, None)
            log.info('azure usage cache loaded: %s', self.azure_usage)

    # ---------- 基础 ----------
    @property
    def _config_path(self):
        return os.environ.get('TRAFFIC_CONFIG', '/home/agent/config.json')

    def _load_all_months(self):
        root = self.config['data_dir']
        if not os.path.isdir(root):
            return
        for year_dir in os.listdir(root):
            ypath = os.path.join(root, year_dir)
            if not os.path.isdir(ypath):
                continue
            for fn in os.listdir(ypath):
                if not fn.endswith('.json'):
                    continue
                try:
                    with open(os.path.join(ypath, fn), 'r', encoding='utf-8') as f:
                        m = json.load(f)
                    mk = f"{m['year']:04d}-{m['month']:02d}"
                    self.months[mk] = m
                except Exception as e:
                    log.warning('load %s/%s failed: %s', ypath, fn, e)

    def _load_alert_state(self):
        self.alert_state = load_json(
            self.config['alert_state_path'],
            {'alert_thresholds_sent': {}, 'monthly_reports_sent': []})
        self.alert_state.setdefault('alert_thresholds_sent', {})
        self.alert_state.setdefault('monthly_reports_sent', [])

    def _save_alert_state(self):
        save_json(self.config['alert_state_path'], self.alert_state)

    def _month_path(self, year, month):
        return os.path.join(self.config['data_dir'], str(year), f'{year:04d}-{month:02d}.json')

    def _persist_month(self, mk):
        m = self.months.get(mk)
        if m is None:
            return
        path = self._month_path(m['year'], m['month'])
        os.makedirs(os.path.dirname(path), exist_ok=True)
        save_json(path, m)

    # ---------- vnstat 读取 ----------
    def _vnstat(self, what):
        bin_path = self.config.get('vnstat_bin', 'vnstat')
        iface = self.config.get('interface', 'eth0')
        cmd = [bin_path, '-i', iface, '--json', what]
        try:
            out = subprocess.run(cmd, capture_output=True, timeout=20, check=True)
            return json.loads(out.stdout.decode('utf-8'))
        except FileNotFoundError:
            log.error('vnstat binary not found: %s', bin_path)
            return None
        except subprocess.TimeoutExpired:
            log.error('vnstat --json %s timed out', what)
            return None
        except (subprocess.CalledProcessError, json.JSONDecodeError, OSError) as e:
            log.error('vnstat --json %s failed: %s', what, e)
            return None

    def poll_vnstat(self):
        """拉取 vnstat 月/日数据，更新并落盘。返回更新的月数。"""
        monthly = self._vnstat('m')
        daily = self._vnstat('d')
        touched = set()

        with self.lock:
            if monthly:
                try:
                    for itf in monthly.get('interfaces', []):
                        for item in itf.get('traffic', {}).get('month', []):
                            d = item.get('date', {})
                            mk = f"{int(d['year']):04d}-{int(d['month']):02d}"
                            m = self.months.get(mk)
                            if m is None:
                                m = month_struct(d['year'], d['month'])
                                self.months[mk] = m
                            m['rx_bytes'] = int(item.get('rx', 0))
                            m['tx_bytes'] = int(item.get('tx', 0))
                            m['updated_at'] = utcnow().isoformat()
                            touched.add(mk)
                except Exception as e:
                    log.exception('parse vnstat monthly failed: %s', e)

            if daily:
                try:
                    by_month = {}
                    for itf in daily.get('interfaces', []):
                        for item in itf.get('traffic', {}).get('day', []):
                            d = item.get('date', {})
                            mk = f"{int(d['year']):04d}-{int(d['month']):02d}"
                            by_month.setdefault(mk, []).append({
                                'day': int(d['day']),
                                'rx_bytes': int(item.get('rx', 0)),
                                'tx_bytes': int(item.get('tx', 0)),
                            })
                    for mk, days in by_month.items():
                        m = self.months.get(mk)
                        if m is None:
                            y, mon = mk.split('-')
                            m = month_struct(int(y), int(mon))
                            self.months[mk] = m
                        days.sort(key=lambda x: x['day'])
                        m['days'] = days
                        m['updated_at'] = utcnow().isoformat()
                        touched.add(mk)
                except Exception as e:
                    log.exception('parse vnstat daily failed: %s', e)

            for mk in touched:
                self._persist_month(mk)
        return len(touched)

    # ---------- 用量（按 Azure 计费周期） ----------
    def current_month_key(self):
        return utcnow().strftime('%Y-%m')

    def cycle_start(self, now):
        """返回当前计费周期的起始 datetime（每月 billing_cycle_start_day 00:00 UTC）。"""
        day = int(self.config.get('billing_cycle_start_day', 1))
        day = max(1, min(28, day))
        if now.day >= day:
            return now.replace(day=day, hour=0, minute=0, second=0, microsecond=0)
        prev = (now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
                - timedelta(days=1))
        return prev.replace(day=day, hour=0, minute=0, second=0, microsecond=0)

    def cycle_end(self, start):
        """周期结束 = 下个月同一天的 00:00（不含）。"""
        nxt = (start.replace(day=1) + timedelta(days=32)).replace(day=1)
        return nxt.replace(day=start.day, hour=0, minute=0, second=0, microsecond=0)

    def _day_totals(self):
        """返回 [{date, rx, tx}]，date 为 UTC date。"""
        out = []
        for m in self.months.values():
            for d in m.get('days', []):
                out.append((date(m['year'], m['month'], d['day']),
                            int(d.get('rx_bytes', 0)), int(d.get('tx_bytes', 0))))
        return out

    def cycle_usage(self):
        """返回当前计费周期 {rx_bytes, tx_bytes, start, end}。"""
        now = utcnow()
        start = self.cycle_start(now)
        end = self.cycle_end(start)
        rx = tx = 0
        with self.lock:
            for dt, r, t in self._day_totals():
                if start.date() <= dt <= now.date():
                    rx += r
                    tx += t
        return {'rx_bytes': rx, 'tx_bytes': tx, 'start': start, 'end': end}

    def tx_bytes_for(self, mk):
        with self.lock:
            m = self.months.get(mk, {})
            return int(m.get('tx_bytes', 0))

    def current_status(self):
        """当前状态：以 Azure 门户口径（Standard Data Transfer Out - Free）为主计量。"""
        now = utcnow()
        quota_gb = float(self.config.get('azure_free_limit_gb', 15))
        cu = self.cycle_usage()
        az = self.azure_usage or {}
        used_gb = az.get('portal_data_transfer_out_gb')
        used_pct = round(used_gb / quota_gb * 100, 4) if (used_gb is not None and quota_gb) else None
        with self.lock:
            cal_mk = now.strftime('%Y-%m')
            m = self.months.get(cal_mk)
            cal_tx = int(m.get('tx_bytes', 0)) if m else 0
            cal_rx = int(m.get('rx_bytes', 0)) if m else 0
            vnstat_updated = (m or {}).get('updated_at')
        return {
            'year': now.year,
            'month': now.month,
            'meter': 'azure_portal',
            'used_gb': used_gb,
            'quota_gb': quota_gb,
            'quota_used_pct': used_pct,
            'billing_cycle_start_day': int(self.config.get('billing_cycle_start_day', 1)),
            'cycle_start': cu['start'].strftime('%Y-%m-%d'),
            'cycle_end': cu['end'].strftime('%Y-%m-%d'),
            'azure_portal_data_transfer_out_gb': az.get('portal_data_transfer_out_gb'),
            'azure_bandwidth_out_gb': az.get('bandwidth_out_gb'),
            'azure_usage_cost_cny': az.get('cost_cny'),
            'azure_usage_month': az.get('month'),
            'azure_usage_updated_at': az.get('updated_at'),
            'vnstat_calendar_month_tx_bytes': cal_tx,
            'vnstat_calendar_month_rx_bytes': cal_rx,
            'vnstat_updated_at': vnstat_updated,
            'updated_at': now.isoformat(),
        }

    def _vnstat_calendar_tx_gb(self):
        """参考值：vnstat 自然月出站（GB，十进制）。"""
        now = utcnow()
        with self.lock:
            m = self.months.get(now.strftime('%Y-%m'), {})
            return int(m.get('tx_bytes', 0)) / GB

    # ---------- Azure 计费用量（Usage Details API） ----------
    def _azure_token(self):
        """通过 VM 托管身份从 IMDS 获取 management.azure.com 的访问令牌。"""
        url = ('http://169.254.169.254/metadata/identity/oauth2/token'
               '?api-version=2018-02-01&resource=https%3A%2F%2Fmanagement.azure.com%2F')
        req = urllib.request.Request(url, headers={'Metadata': 'true'})
        with urllib.request.urlopen(req, timeout=15) as resp:
            return json.loads(resp.read().decode('utf-8'))['access_token']

    def _azure_subscription_id(self):
        url = 'http://169.254.169.254/metadata/instance/compute?api-version=2021-02-01'
        req = urllib.request.Request(url, headers={'Metadata': 'true'})
        with urllib.request.urlopen(req, timeout=15) as resp:
            return json.loads(resp.read().decode('utf-8')).get('subscriptionId')

    def fetch_azure_usage(self, month_offset=0):
        """拉取指定月份 usageDetails，聚合与门户免费服务页同源的 Data Transfer Out 用量。
        month_offset: 0=当月（写入缓存），-1=上月（用于月报，不覆盖缓存）。"""
        sub = self._azure_subscription_id()
        if not sub:
            raise RuntimeError('cannot read subscriptionId from IMDS')
        token = self._azure_token()
        now = utcnow()
        month_start = (now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
                       + timedelta(days=month_offset * 32)).replace(day=1)
        nxt = (month_start.replace(day=1) + timedelta(days=32)).replace(day=1)
        flt = ("properties/usageStart ge '{0}T00:00:00Z' and "
               "properties/usageStart lt '{1}T00:00:00Z'").format(
                   month_start.strftime('%Y-%m-%d'), nxt.strftime('%Y-%m-%d'))
        base = ('https://management.azure.com/subscriptions/{sub}/providers/'
                'Microsoft.Consumption/usageDetails?api-version=2023-11-01'
                '&metric=actualcost&$filter={f}')
        url = base.format(sub=sub, f=urllib.parse.quote(flt, safe=''))
        portal_gb = 0.0
        bw_out_gb = 0.0
        cost_cny = 0.0
        rows = 0
        pages = 0
        while url:
            req = urllib.request.Request(url, headers={'Authorization': 'Bearer ' + token})
            with urllib.request.urlopen(req, timeout=30) as resp:
                d = json.loads(resp.read().decode('utf-8'))
            for item in d.get('value', []):
                p = item.get('properties', {}) or {}
                qty = float(p.get('quantity') or 0)
                rows += 1
                name = str(p.get('meterName') or '')
                subcat = str(p.get('meterSubCategory') or '')
                if str(p.get('meterCategory') or '') == 'Bandwidth' and 'Data Transfer Out' in name:
                    bw_out_gb += qty
                    # 门户免费服务页 "Data Transfer Out (GB)" = MGN 路由的出站免费计费表
                    if 'MGN' in subcat or '- Free' in name:
                        portal_gb += qty
                    cost_cny += float(p.get('costInBillingCurrency') or 0)
            pages += 1
            url = d.get('nextLink')
            if pages > 30:
                log.warning('azure usage pagination limit reached')
                break
        result = {
            'month': month_start.strftime('%Y-%m'),
            'portal_data_transfer_out_gb': round(portal_gb, 4),
            'bandwidth_out_gb': round(bw_out_gb, 4),
            'cost_cny': round(cost_cny, 4),
            'rows': rows,
            'updated_at': utcnow().isoformat(),
        }
        self.azure_usage = result
        if month_offset == 0:
            path = self.config.get('azure_usage_cache_path')
            if path:
                save_json(path, result)
        log.info('azure usage fetched (offset=%d): %s', month_offset, result)
        return result

    def check_azure_usage(self):
        if not self.config.get('azure_usage_enabled', False):
            return
        interval = float(self.config.get('azure_usage_fetch_interval_sec', 21600))
        now = time.time()
        if now - self._azure_last_fetch < interval:
            return
        self._azure_last_fetch = now
        try:
            self.fetch_azure_usage()
        except Exception as e:
            log.warning('azure usage fetch failed: %s', e)

    # ---------- IP / 邮件 ----------
    def get_current_ip(self):
        configured = self.config.get('current_ip')
        if configured:
            return str(configured)
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try:
                s.connect(('8.8.8.8', 80))
                return s.getsockname()[0]
            finally:
                s.close()
        except OSError:
            return '0.0.0.0'

    def send_email(self, subject, body):
        cfg = self.config['smtp']
        msg = EmailMessage()
        msg['Subject'] = subject
        msg['From'] = formataddr((cfg.get('from_name', ''), cfg.get('from_addr')))
        msg['To'] = cfg.get('to_addr')
        msg.set_content(body)

        context = None
        if cfg.get('skip_cert_verify'):
            context = ssl._create_unverified_context()
        host = cfg['host']
        port = int(cfg.get('port', 465))
        timeout = int(cfg.get('timeout_sec', 30))
        if cfg.get('ssl', False):
            smtp_cls = smtplib.SMTP_SSL
            kwargs = {'context': context}
        else:
            smtp_cls = smtplib.SMTP
            kwargs = {}
        with smtp_cls(host, port, timeout=timeout, **kwargs) as smtp:
            if not cfg.get('ssl', False):
                smtp.ehlo()
                if cfg.get('starttls', True):
                    smtp.starttls(context=context)
                    smtp.ehlo()
            smtp.login(cfg['user'], cfg['password'])
            smtp.send_message(msg)
        log.info('email sent: %s', subject)

    def build_alert_mail(self, mk, used_gb, quota_gb, pct, az):
        ip = self.get_current_ip()
        month = (az or {}).get('month', mk.replace('azure:', ''))
        az_updated = (az or {}).get('updated_at', '')
        vnstat_tx = self._vnstat_calendar_tx_gb()
        body = (
            'Royirin Rosmontis Private Proxy JP 流量监控\n'
            '==========\n'
            '计量口径：Azure 门户（Standard Data Transfer Out - Free）\n'
            f'统计月份：{month}（自然月）\n'
            f'当前用量：{used_gb:.2f} GB / {quota_gb:.0f} GB（{pct:.2f}%）\n'
            f'AZURE 数据更新：{az_updated}\n'
            f'参考(vnstat 自然月出站)：{vnstat_tx:.2f} GB\n'
            f'当前 IP：{ip}\n'
            f'生效域名：{self.config.get("effective_domain", "")}\n'
            '==========\n'
            f'已达阈值 {pct:.2f}% (Azure 门户口径)，请注意使用情况'
        )
        subject = f'[流量告警] {month} 已用 {pct:.2f}% ({used_gb:.2f}/{quota_gb:.0f} GB)'
        return subject, body

    def build_monthly_mail(self, month, used_gb, quota_gb, pct, az):
        ip = self.get_current_ip()
        over = '是' if used_gb > quota_gb else '否'
        az_updated = (az or {}).get('updated_at', '')
        vnstat_tx = self._vnstat_calendar_tx_gb()
        body = (
            'Royirin Rosmontis Private Proxy JP 流量监控\n'
            '月度报告\n'
            '==========\n'
            '计量口径：Azure 门户（Standard Data Transfer Out - Free）\n'
            f'报告月份：{month}（自然月）\n'
            f'当月用量：{used_gb:.2f} GB / {quota_gb:.0f} GB（{pct:.2f}%）\n'
            f'AZURE 数据更新：{az_updated}\n'
            f'参考(vnstat 自然月出站)：{vnstat_tx:.2f} GB\n'
            f'当前 IP：{ip}\n'
            f'生效域名：{self.config.get("effective_domain", "")}\n'
            '==========\n'
            f'上月使用 {over} 超限, 使用量为限额的 {pct:.2f}%'
        )
        subject = f'[流量月报] {month} 用量 {used_gb:.2f} GB ({pct:.2f}%)'
        return subject, body

    # ---------- 定时任务 ----------
    def check_alerts(self):
        """阈值告警：以 Azure 门户口径为准（无 Azure 数据时跳过，避免误报）。"""
        quota_gb = float(self.config.get('azure_free_limit_gb', 15))
        if quota_gb <= 0:
            return
        az = self.azure_usage or {}
        used_gb = az.get('portal_data_transfer_out_gb')
        if used_gb is None:
            log.warning('alerts skipped: no azure usage data yet')
            return
        mk = 'azure:' + str(az.get('month') or utcnow().strftime('%Y-%m'))
        pct = used_gb / quota_gb * 100
        sent = self.alert_state['alert_thresholds_sent'].setdefault(mk, [])
        if not sent:
            # 新周期/口径切换：已越过的档位直接标记，避免切换瞬间补发一堆邮件
            for t in sorted(self.config.get('alert_thresholds_pct', [])):
                if pct >= t:
                    sent.append(t)
            self._save_alert_state()
            log.info('seeded thresholds for %s at %.2f%%: %s', mk, pct, sent)
        for t in sorted(self.config.get('alert_thresholds_pct', [])):
            if pct >= t and t not in sent:
                subject, body = self.build_alert_mail(mk, used_gb, quota_gb, pct, az)
                try:
                    self.send_email(subject, body)
                    sent.append(t)
                    self._save_alert_state()
                    log.info('threshold %s%% alert sent for %s (used %.2f%%)', t, mk, pct)
                except Exception as e:
                    log.error('send threshold %s%% alert failed for %s: %s; will retry',
                              t, mk, e)
                    break

    def check_monthly_report(self):
        """月度报告：以 Azure 门户口径为准（上月数据，当月 1 号后发送）。"""
        now = utcnow()
        cur_start = self.cycle_start(now)
        prev_start = (cur_start.replace(day=1) - timedelta(days=1)).replace(day=1)
        prev_mk = 'azure:' + prev_start.strftime('%Y-%m')
        if prev_mk in self.alert_state['monthly_reports_sent']:
            return
        # 口径切换迁移：旧键(cycle:)已发过则直接标记，避免重复发信
        old_mk = 'cycle:' + prev_start.strftime('%Y-%m')
        if old_mk in self.alert_state['monthly_reports_sent']:
            self.alert_state['monthly_reports_sent'].append(prev_mk)
            self._save_alert_state()
            log.info('monthly report key migrated: %s already sent as %s', prev_mk, old_mk)
            return
        try:
            az_prev = self.fetch_azure_usage(month_offset=-1)
        except Exception as e:
            log.error('fetch prev month azure usage failed: %s; will retry', e)
            return
        quota_gb = float(self.config.get('azure_free_limit_gb', 15))
        used_gb = az_prev.get('portal_data_transfer_out_gb') or 0.0
        pct = used_gb / quota_gb * 100 if quota_gb else 0.0
        subject, body = self.build_monthly_mail(
            az_prev.get('month', prev_start.strftime('%Y-%m')), used_gb, quota_gb, pct, az_prev)
        try:
            self.send_email(subject, body)
            self.alert_state['monthly_reports_sent'].append(prev_mk)
            self._save_alert_state()
            log.info('monthly report sent for %s (used %.2f GB, %.2f%%)',
                     prev_mk, used_gb, pct)
        except Exception as e:
            log.error('send monthly report for %s failed: %s; will retry', prev_mk, e)


def build_http_handler(monitor):
    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):
            if self.path == '/vpn/status/traffic':
                try:
                    payload = json.dumps(monitor.current_status(), ensure_ascii=False)
                    body = payload.encode('utf-8')
                    self.send_response(200)
                    self.send_header('Content-Type', 'application/json; charset=utf-8')
                    self.send_header('Content-Length', str(len(body)))
                    self.end_headers()
                    self.wfile.write(body)
                except Exception as e:
                    log.exception('status endpoint error')
                    body = b'{"error":"internal"}'
                    self.send_response(500)
                    self.send_header('Content-Type', 'application/json')
                    self.send_header('Content-Length', str(len(body)))
                    self.end_headers()
                    self.wfile.write(body)
            elif self.path in ('/', '/healthz'):
                body = b'ok\n'
                self.send_response(200)
                self.send_header('Content-Type', 'text/plain')
                self.send_header('Content-Length', str(len(body)))
                self.end_headers()
                self.wfile.write(body)
            else:
                body = b'{"error":"not found"}\n'
                self.send_response(404)
                self.send_header('Content-Type', 'application/json')
                self.send_header('Content-Length', str(len(body)))
                self.end_headers()
                self.wfile.write(body)

        def log_message(self, fmt, *args):
            log.debug('http: ' + fmt, *args)

    return Handler


def run_service(config):
    monitor = Monitor(config)

    def poll_loop():
        while True:
            try:
                n = monitor.poll_vnstat()
                log.debug('vnstat poll updated %d month(s)', n)
            except Exception:
                log.exception('poll loop error')
            time.sleep(int(config.get('poll_interval_sec', 60)))

    def schedule_loop():
        while True:
            try:
                monitor.check_azure_usage()  # 先刷新 Azure 口径数据，再评估告警
                monitor.check_alerts()
                monitor.check_monthly_report()
            except Exception:
                log.exception('schedule loop error')
            time.sleep(int(config.get('schedule_interval_sec', 30)))

    threading.Thread(target=poll_loop, daemon=True, name='vnstat-poll').start()
    threading.Thread(target=schedule_loop, daemon=True, name='scheduler').start()

    host = config.get('http_host', '127.0.0.1')
    port = int(config.get('http_port', 18080))
    server = ThreadingHTTPServer((host, port), build_http_handler(monitor))
    log.info('HTTP server listening on http://%s:%d/vpn/status/traffic', host, port)
    server.serve_forever()


def main():
    ap = argparse.ArgumentParser(description='vnstat traffic monitor')
    ap.add_argument('--config', default=os.environ.get('TRAFFIC_CONFIG', '/home/agent/config.json'))
    ap.add_argument('--once', action='store_true', help='拉取一次 vnstat 后退出')
    ap.add_argument('--status', action='store_true', help='打印当前状态 JSON 后退出')
    ap.add_argument('--test-email', action='store_true', help='发送一封测试邮件后退出')
    args = ap.parse_args()

    os.environ['TRAFFIC_CONFIG'] = args.config
    with open(args.config, 'r', encoding='utf-8') as f:
        config = json.load(f)

    logging.basicConfig(
        level=getattr(logging, config.get('log_level', 'INFO').upper(), logging.INFO),
        format='%(asctime)s %(levelname)s %(name)s: %(message)s')

    if args.once:
        m = Monitor(config)
        n = m.poll_vnstat()
        print(f'updated {n} month(s)')
        return 0
    if args.status:
        m = Monitor(config)
        m.poll_vnstat()
        print(json.dumps(m.current_status(), ensure_ascii=False, indent=2))
        return 0
    if args.test_email:
        m = Monitor(config)
        m.poll_vnstat()
        quota_gb = float(config.get('azure_free_limit_gb', 15))
        az = m.azure_usage or {}
        used_gb = az.get('portal_data_transfer_out_gb')
        if used_gb is None:
            try:
                m.fetch_azure_usage()
                az = m.azure_usage or {}
                used_gb = az.get('portal_data_transfer_out_gb')
            except Exception as e:
                print('azure usage fetch failed:', e)
        mk = 'azure:' + str(az.get('month') or utcnow().strftime('%Y-%m'))
        pct = used_gb / quota_gb * 100 if (used_gb is not None and quota_gb) else 0.0
        subject, body = m.build_alert_mail(mk, used_gb or 0.0, quota_gb, pct, az)
        print('sending test email:')
        print('Subject:', subject)
        print(body)
        m.send_email(subject, body)
        print('sent')
        return 0
    run_service(config)


if __name__ == '__main__':
    sys.exit(main())