#!/usr/bin/env python3
"""量測「這個儲存裝置到底能有多快」——本專案所有 I/O 最佳化的前提。

為什麼需要這個工具：分頁引擎的每一層都要等 SSD。如果不知道單次讀取的延遲與
頻寬，就只能靠猜（例如「多開執行緒會更快」）。實測結果常常相反：
**單一 1.9 MB 隨機讀約 7 ms，而且執行緒再多也不會變快**（見下方輸出），
因為瓶頸是單流頻寬（約 275 MB/s），不是延遲也不是佇列深度。這正是
「I/O 段平行」這個最佳化能生效的原因：它不是開更多執行緒，而是把
「一個 expert 的三段循序讀」變成「同時讀」。

用法
----
    python tools/io_probe.py --model /path/to/model.gguf
    python tools/io_probe.py --model xxx.gguf --chunk 1900544 --parallel 1,2,4,8,16,24

輸出
----
每個併發度一列：median / mean / p90 的單次讀取延遲，以及整體聚合頻寬。
"""

from __future__ import annotations

import argparse
import os
import random
import statistics
import threading
import time

PAGE = os.sysconf("SC_PAGE_SIZE")


def drop_cache(fd: int, off: int, length: int) -> None:
    """把剛讀過的範圍從 kernel page cache 丟掉，否則下一輪量到的是 RAM 不是 SSD。"""
    start = off - (off % PAGE)
    end = ((off + length + PAGE - 1) // PAGE) * PAGE
    try:
        os.posix_fadvise(fd, start, end - start, os.POSIX_FADV_DONTNEED)
    except OSError:
        pass


def measure(fd: int, size: int, chunk: int, parallel: int, n: int) -> dict:
    random.seed(11)
    max_off = max(1, size - chunk - PAGE)
    offsets = [random.randrange(0, max_off) for _ in range(n)]
    bufs = [bytearray(chunk) for _ in range(parallel)]
    idx = [0]
    lock = threading.Lock()
    lat: list[float] = []

    def work(tid: int) -> None:
        local: list[float] = []
        while True:
            with lock:
                i = idx[0]
                if i >= len(offsets):
                    break
                idx[0] += 1
            off = offsets[i]
            mv = memoryview(bufs[tid])
            t0 = time.perf_counter()
            os.preadv(fd, [mv], off)
            local.append((time.perf_counter() - t0) * 1000.0)
            drop_cache(fd, off, chunk)
        with lock:
            lat.extend(local)

    ths = [threading.Thread(target=work, args=(i,)) for i in range(parallel)]
    t_start = time.perf_counter()
    for t in ths:
        t.start()
    for t in ths:
        t.join()
    wall = time.perf_counter() - t_start
    lat.sort()
    return {
        "parallel": parallel,
        "n": len(lat),
        "median_ms": round(statistics.median(lat), 2),
        "mean_ms": round(statistics.mean(lat), 2),
        "p90_ms": round(lat[int(len(lat) * 0.9)], 2),
        "wall_s": round(wall, 3),
        "agg_MBps": round(len(lat) * chunk / wall / 1e6, 1),
    }


def main() -> int:
    ap = argparse.ArgumentParser(description="SDQ 儲存裝置 I/O 特性量測")
    ap.add_argument("--model", required=True, help="拿來當讀取目標的檔案（通常就是 gguf）")
    ap.add_argument("--chunk", type=int, default=1_900_000, help="單次讀取的位元組數")
    ap.add_argument("--parallel", default="1,2,4,8,16,24", help="要量哪幾種併發度")
    ap.add_argument("--n", type=int, default=64, help="每個併發度讀幾次")
    args = ap.parse_args()

    size = os.path.getsize(args.model)
    fd = os.open(args.model, os.O_RDONLY)
    print(f"檔案：{args.model}（{size / 2**30:.1f} GiB）")
    print(f"單次讀取 {args.chunk} B，每個併發度 {args.n} 次")
    print()
    print(f"{'併發':>5} {'中位延遲':>10} {'平均':>8} {'p90':>8} {'聚合頻寬':>12}")
    try:
        for p in [int(x) for x in args.parallel.split(",")]:
            r = measure(fd, size, args.chunk, p, args.n)
            print(f"{r['parallel']:5d} {r['median_ms']:9.2f}ms {r['mean_ms']:7.2f}ms "
                  f"{r['p90_ms']:7.2f}ms {r['agg_MBps']:9.1f} MB/s")
    finally:
        os.close(fd)
    print()
    print("讀法：如果『中位延遲』不隨併發度下降，代表單次讀取是頻寬受限的；")
    print("     這時候開更多執行緒只會讓聚合頻寬變高，不會縮短關鍵路徑。")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())