#!/usr/bin/env python3
"""SDQ benchmark: measure pager/I/O/cache/tok/s and report RSS/VmSwap.

這個模組是**舊的**基準工具，它給的數字不是 8 GiB / no-swap 證據：

1. 預設的 `RLIMIT_DATA` 只限制 data/heap 配置，不計 file-backed 常駐頁，
   也不是 swap 開關。
2. 報告的 `VmSwap` 只是「那一次執行沒有被換出」，不是「核心不可能換出」。

**要 8 GiB / no-swap 的證據請用 `tools/strict8g.py` 或 `sdq.strict_bench.py`。**
本模組保留給快速的非受管治量測，以及 `demo.py --bench` 呼叫。

做法
----
1. 每次量測前用 tools/drop_model_cache.py 把模型檔的 page cache 丟掉，
   避免「上一輪留在核心 RAM 的頁」讓 SSD 讀取量看起來比實際少。

用法
----
    python -m sdq.bench --model /path/model.gguf --bin build/bin/epq-chat
    python -m sdq.bench --ram-mb 8192 --arena-mb 3800 --threads 4

受管治的版本（會輸出 strict_ok）::

    python -m sdq.strict_bench --model /path/model.gguf --bin build/bin/epq-chat \
        --rounds 9 --strict-mb 8192 --label headline
"""

from __future__ import annotations

import argparse
import json
import os
import resource
import subprocess
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
PROMPTS = [
    "請用三句話說明什麼是 SSD、RAM、CPU 之間的資料分層。",
    "Write a short paragraph about how mixture-of-experts models route tokens.",
    "請列出五種日常會用到電腦的工作，並各用一句話說明。",
]


def drop_cache(model: str) -> None:
    subprocess.run([sys.executable, str(REPO / "tools" / "drop_model_cache.py"), model],
                   check=False, capture_output=True)


def run_once(binary: str, model: str, args: argparse.Namespace, prompt: str,
             n_predict: int, tag: str) -> dict:
    stats_file = Path(args.workdir) / f"stats_{tag}.jsonl"
    if stats_file.exists():
        stats_file.unlink()

    cmd = [
        binary, "-m", model, "-p", prompt, "-n", str(n_predict),
        "-c", str(args.ctx), "-t", str(args.threads), "-b", str(args.ubatch),
        "--temp", "0", "--no-chat", "--no-interactive",
        "--stats-file", str(stats_file),
    ]
    if args.ram_mb:
        cmd += ["--ram-budget-mb", str(args.ram_mb)]
    if args.arena_mb:
        cmd += ["--arena-mb", str(args.arena_mb)]
    if args.no_prefetch:
        cmd += ["--no-prefetch"]

    env = dict(os.environ)
    env.setdefault("SDQ_FADV_DONTNEED", "1")

    if args.ram_mb and not args.no_ulimit:
        # RLIMIT_DATA limits data/heap allocations, not total RSS or swap.
        def _limit():
            limit = args.ram_mb * 1024 * 1024
            resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
        preexec = _limit
    else:
        preexec = None

    if args.drop_cache:
        drop_cache(model)

    t0 = time.time()
    proc = subprocess.run(cmd, capture_output=True, text=True, env=env, preexec_fn=preexec)
    wall = time.time() - t0
    if proc.returncode != 0:
        print(f"[bench] {tag} 執行失敗（rc={proc.returncode}）")
        print(proc.stderr[-2000:])
        return {"tag": tag, "error": proc.returncode}

    # 解析指標
    stats = {}
    if stats_file.exists():
        for line in stats_file.read_text().splitlines():
            rec = json.loads(line)
            stats[rec["phase"]] = rec

    mem = {}
    for line in proc.stderr.splitlines():
        if "記憶體（載入後）" in line or "記憶體：" in line:
            mem["line"] = line
            for key, val in re_findall(line):
                mem[key] = val

    tok = {}
    for line in proc.stderr.splitlines():
        if "decode:" in line:
            tok["decode"] = line
        if "prefill:" in line:
            tok["prefill"] = line

    return {
        "tag": tag,
        "wall_s": round(wall, 2),
        "stdout": proc.stdout.strip()[-400:],
        "tok": tok,
        "mem": mem,
        "stats": stats,
    }


def re_findall(line: str):
    import re
    return re.findall(r"(RSS|VmSwap|峰值 RSS\(VmHWM\))\)?[:=]?\s*(\d+)\s*MB", line)


def summarize(result: dict) -> str:
    s = result["stats"].get("decode") or result["stats"].get("bench_hot") or {}
    if not s:
        return f"[{result['tag']}] 無指標"
    return (
        f"[{result['tag']}] hit={s['hit_rate']:.1f}%  SSD/token={s['ssd_bytes_per_token']/2**20:.1f} MB  "
        f"預取使用率={s['prefetch_used_rate']:.1f}%  arena={s['arena_used_bytes']/2**30:.2f}/{s['arena_bytes']/2**30:.2f} GiB  "
        f"RSS={s['rss_bytes']/2**20:.0f} MB"
    )


def main() -> None:
    ap = argparse.ArgumentParser(description="SDQ 8GB RAM 基準測量")
    ap.add_argument("--model", required=True)
    ap.add_argument("--bin", default=str(REPO / "build" / "bin" / "epq-chat"))
    ap.add_argument("--ram-mb", type=int, default=8192,
                    help="SDQ pager sizing budget (also RLIMIT_DATA unless --no-ulimit; not an RSS/swap cap)")
    ap.add_argument("--arena-mb", type=int, default=0, help="固定 arena 大小（0 = 由預算推算）")
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--ubatch", type=int, default=128)
    ap.add_argument("--ctx", type=int, default=4096)
    ap.add_argument("--n-predict", type=int, default=48)
    ap.add_argument("--rounds", type=int, default=2)
    ap.add_argument("--workdir", default=str(REPO / "bench" / "results"))
    ap.add_argument("--no-prefetch", action="store_true")
    ap.add_argument("--no-ulimit", action="store_true",
                    help="Do not set RLIMIT_DATA; still not an OS memory/swap cap")
    ap.add_argument("--no-drop-cache", dest="drop_cache", action="store_false", default=True)
    ap.add_argument("--label", default="run")
    args = ap.parse_args()

    Path(args.workdir).mkdir(parents=True, exist_ok=True)

    results = []
    # 冷啟動：整個模型第一次被讀（cache 是空的）
    results.append(run_once(args.bin, args.model, args, PROMPTS[0], args.n_predict, "cold"))
    print(summarize(results[-1]))
    # 熱身後：專家 cache 已經有內容
    for r in range(args.rounds):
        res = run_once(args.bin, args.model, args, PROMPTS[(r + 1) % len(PROMPTS)],
                       args.n_predict, f"hot{r}")
        results.append(res)
        print(summarize(res))

    out = Path(args.workdir) / f"bench_{args.label}.json"
    out.write_text(json.dumps({
        "config": vars(args),
        "results": results,
    }, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"\n[bench] 完整結果已寫到 {out}")

    # 印出對照用的解譯文字（確認模型真的有在回答，不是亂碼）
    for res in results:
        print(f"\n--- {res['tag']} 生成內容 ---\n{res.get('stdout','')}")


if __name__ == "__main__":
    main()