#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业15 · Worker 宕机实验：观察 Redis 中 Celery 队列长度

在隔离环境停 worker 前后各跑一次，对比队列是否堆积。

用法：
  python hw15_02_观察Redis队列.py
  python hw15_02_观察Redis队列.py --redis-host 127.0.0.1 --redis-port 6380
  python hw15_02_观察Redis队列.py --label after_worker_stop

输出：
  data/homework15/redis_queue_snapshot_<label>_<timestamp>.json
  终端打印各队列 LLEN
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path

# Dify 常见 Celery 队列名（与 worker 消费一致）
DEFAULT_QUEUES = [
    "dataset",
    "mail",
    "ops_trace",
    "app_deletion",
    "workflow",
    "celery",  # 默认队列
]


def out_dir() -> Path:
    env = os.environ.get("HW15_OUT")
    if env:
        p = Path(env)
        p.mkdir(parents=True, exist_ok=True)
        return p
    repo = Path(__file__).resolve().parents[2] / "data" / "homework15"
    repo.mkdir(parents=True, exist_ok=True)
    return repo


def redis_cli(args: list[str]) -> str:
    try:
        r = subprocess.run(args, capture_output=True, text=True, timeout=15)
    except FileNotFoundError:
        print("redis-cli 未找到；请安装 redis-tools 或在容器内执行", file=sys.stderr)
        sys.exit(2)
    if r.returncode != 0:
        print(r.stderr or r.stdout, file=sys.stderr)
        sys.exit(r.returncode)
    return (r.stdout or "").strip()


def queue_key(queue: str) -> str:
    # Celery Redis broker 默认 list key
    return queue


def main() -> None:
    parser = argparse.ArgumentParser(description="观察 Celery Redis 队列长度")
    parser.add_argument("--redis-host", default=os.environ.get("REDIS_HOST", "127.0.0.1"))
    parser.add_argument("--redis-port", type=int, default=int(os.environ.get("REDIS_PORT", "6379")))
    parser.add_argument("--redis-password", default=os.environ.get("REDIS_PASSWORD", "difyai123456"))
    parser.add_argument("--label", default="snapshot", help="快照标签，如 baseline / after_worker_stop")
    parser.add_argument("--queues", nargs="*", default=DEFAULT_QUEUES)
    args = parser.parse_args()

    base = [
        "redis-cli",
        "-h",
        args.redis_host,
        "-p",
        str(args.redis_port),
        "-a",
        args.redis_password,
        "--no-auth-warning",
    ]

    ping = redis_cli(base + ["PING"])
    if ping.upper() != "PONG":
        print(f"Redis PING 异常: {ping}", file=sys.stderr)
        sys.exit(1)

    snapshot: dict = {
        "at": datetime.now().isoformat(timespec="seconds"),
        "label": args.label,
        "redis": f"{args.redis_host}:{args.redis_port}",
        "queues": {},
    }

    print(f"Redis {args.redis_host}:{args.redis_port}  label={args.label}")
    print("-" * 48)
    for q in args.queues:
        key = queue_key(q)
        try:
            n = int(redis_cli(base + ["LLEN", key]))
        except ValueError:
            n = -1
        snapshot["queues"][q] = {"key": key, "llen": n}
        print(f"  {q:16s}  LLEN={n}")

    # 额外：未消费任务总数粗估（keys 仅作实验环境参考）
    try:
        keys = redis_cli(base + ["KEYS", "*"])
        snapshot["key_count_hint"] = len([k for k in keys.splitlines() if k.strip()])
    except Exception:
        snapshot["key_count_hint"] = None

    od = out_dir()
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    out = od / f"redis_queue_snapshot_{args.label}_{ts}.json"
    out.write_text(json.dumps(snapshot, ensure_ascii=False, indent=2), encoding="utf-8")
    print("-" * 48)
    print(f"Wrote {out}")


if __name__ == "__main__":
    main()