#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : experiments/micro/expH_ssd_feasibility/benchmark.py
#  Purpose   : Measure real SSD read behavior on macOS/APFS (F_NOCACHE, block
#              sizes, random vs sequential, threads, concurrent Metal load)
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-11
#  Modified  : 2026-08-11
#  Platform  : macOS / Apple Silicon (arm64)
#  License   : All rights reserved (research code)
# =============================================================================
"""Experiment H — SSD feasibility (charter §9.H).

Measures actual (not theoretical) read throughput of the internal Apple NVMe
at block sizes 4 KB … 4 MB, sequential vs random, with and without the page
cache (F_NOCACHE — macOS has no O_DIRECT), 1/4/8 threads, and optionally
under concurrent Metal GPU compute (MLX matmul loop).

Writes results JSON (with embedded hardware manifest) to
results/expH_ssd_feasibility/<timestamp>/results.json.

Usage:
    python3 benchmark.py [--file-gib 8] [--repeats 3] [--quick]
"""

from __future__ import annotations

import argparse
import fcntl
import json
import os
import statistics
import sys
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO_ROOT / "benchmarks"))
from hardware_manifest import collect_manifest  # noqa: E402

F_NOCACHE = 48  # from <sys/fcntl.h> on macOS
BLOCK_SIZES = [4 << 10, 16 << 10, 64 << 10, 256 << 10, 1 << 20, 4 << 20]
THREAD_COUNTS = [1, 4, 8]


def create_test_file(path: Path, size_bytes: int) -> None:
    """Write an incompressible test file WITHOUT populating the page cache.

    Crucial macOS detail (discovered in the first run of this experiment):
    F_NOCACHE on a *read* fd does not bypass pages already resident in the
    unified page cache — and writing the file normally makes every page
    resident. The file must therefore be written with F_NOCACHE set on the
    write fd, so 'cold' reads afterwards genuinely hit the SSD.
    """
    if path.exists():
        path.unlink()  # always recreate: an old file may have cached pages
    import numpy as np

    rng = np.random.default_rng(42)
    chunk = rng.integers(0, 256, size=64 << 20, dtype=np.uint8).tobytes()
    written = 0
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
    try:
        fcntl.fcntl(fd, F_NOCACHE, 1)
        while written < size_bytes:
            n = min(len(chunk), size_bytes - written)
            os.write(fd, chunk[:n])
            written += n
        os.fsync(fd)
    finally:
        os.close(fd)


def _reader(path: Path, offsets: list[int], block: int, nocache: bool,
            out: dict, idx: int) -> None:
    fd = os.open(path, os.O_RDONLY)
    try:
        if nocache:
            fcntl.fcntl(fd, F_NOCACHE, 1)
        total = 0
        t0 = time.perf_counter()
        for off in offsets:
            total += len(os.pread(fd, block, off))
        dt = time.perf_counter() - t0
        out[idx] = (total, dt)
    finally:
        os.close(fd)


def measure(path: Path, file_size: int, block: int, pattern: str,
            nocache: bool, threads: int, budget_bytes: int, seed: int) -> dict:
    """One measurement cell. Returns MB/s and IOPS."""
    import random as _random

    n_blocks_total = max(1, budget_bytes // block)
    max_off = file_size - block
    rng = _random.Random(seed)
    if pattern == "random":
        all_offsets = [rng.randrange(0, max_off // block) * block for _ in range(n_blocks_total)]
    else:
        stride = max(block, (max_off // n_blocks_total) // block * block) if n_blocks_total > 1 else block
        all_offsets = [(i * block) % (max_off + 1) for i in range(n_blocks_total)]
        del stride

    per_thread = [all_offsets[i::threads] for i in range(threads)]
    out: dict = {}
    ts = [
        threading.Thread(target=_reader, args=(path, per_thread[i], block, nocache, out, i))
        for i in range(threads)
    ]
    t0 = time.perf_counter()
    for t in ts:
        t.start()
    for t in ts:
        t.join()
    wall = time.perf_counter() - t0
    total_bytes = sum(v[0] for v in out.values())
    return {
        "mb_per_s": total_bytes / wall / 1e6,
        "iops": total_bytes / block / wall,
        "bytes": total_bytes,
        "wall_s": wall,
    }


def warm_cache(path: Path) -> None:
    with open(path, "rb") as f:
        while f.read(64 << 20):
            pass


def adaptive_budget(path: Path, file_size: int, block: int, pattern: str,
                    nocache: bool, threads: int, floor: int,
                    target_s: float = 3.0) -> int:
    """Choose a per-cell byte budget so each repeat runs ~target_s of wall
    time (short cells are dominated by thread-start overhead and timer
    granularity — the first quick run produced impossible 48 GB/s cells at
    ~10 ms wall)."""
    probe = measure(path, file_size, block, pattern, nocache, threads,
                    budget_bytes=max(floor // 8, 32 << 20), seed=7)
    rate = probe["mb_per_s"] * 1e6
    budget = int(rate * target_s)
    return max(floor, min(budget, 8 << 30))


class IostatLogger:
    """Log `iostat -d -w 1` for the duration of the run — ground truth for
    whether bytes actually came from the disk controller vs the page cache."""

    def __init__(self, out_path: Path) -> None:
        self.out_path = out_path
        self.proc: object = None

    def __enter__(self) -> "IostatLogger":
        import subprocess

        self._fh = open(self.out_path, "w")
        self.proc = subprocess.Popen(
            ["iostat", "-d", "-w", "1"], stdout=self._fh, stderr=subprocess.DEVNULL
        )
        return self

    def __exit__(self, *exc: object) -> None:
        self.proc.terminate()
        self._fh.close()


class GpuLoad:
    """Background MLX matmul loop to contend for unified memory bandwidth."""

    def __init__(self) -> None:
        self.stop = threading.Event()
        self.iterations = 0
        self.thread: threading.Thread | None = None
        self.available = False
        try:
            import mlx.core as mx

            self._mx = mx
            self.available = True
        except ImportError:
            pass

    def _loop(self) -> None:
        mx = self._mx
        a = mx.random.normal((4096, 4096), dtype=mx.float16)
        b = mx.random.normal((4096, 4096), dtype=mx.float16)
        while not self.stop.is_set():
            c = mx.matmul(a, b)
            mx.eval(c)
            self.iterations += 1

    def __enter__(self) -> "GpuLoad":
        if self.available:
            self.thread = threading.Thread(target=self._loop, daemon=True)
            self.thread.start()
            time.sleep(0.5)  # let the GPU ramp
        return self

    def __exit__(self, *exc: object) -> None:
        self.stop.set()
        if self.thread:
            self.thread.join(timeout=10)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--file-gib", type=float, default=8.0)
    ap.add_argument("--repeats", type=int, default=3)
    ap.add_argument("--budget-mib", type=int, default=1024,
                    help="bytes read per measurement cell")
    ap.add_argument("--quick", action="store_true",
                    help="1 repeat, 512 MiB budget, threads {1,8}")
    args = ap.parse_args()
    if args.quick:
        args.repeats, args.budget_mib = 1, 512
        thread_counts = [1, 8]
    else:
        thread_counts = THREAD_COUNTS

    file_size = int(args.file_gib * (1 << 30))
    test_file = Path(__file__).parent / "results" / "testfile.bin"
    test_file.parent.mkdir(exist_ok=True)
    print(f"creating {args.file_gib} GiB test file (once)…", flush=True)
    create_test_file(test_file, file_size)

    manifest = collect_manifest()
    results: list[dict] = []
    budget = args.budget_mib << 20

    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    out_dir = REPO_ROOT / "results" / "expH_ssd_feasibility" / ts
    out_dir.mkdir(parents=True)
    iostat = IostatLogger(out_dir / "iostat.log")
    iostat.__enter__()

    def run_cells(gpu_loaded: bool) -> None:
        for nocache in (True, False):
            if not nocache:
                warm_cache(test_file)
            for pattern in ("random", "sequential"):
                for block in BLOCK_SIZES:
                    for threads in thread_counts:
                        if gpu_loaded and (pattern != "random" or not nocache or threads != 8):
                            continue  # GPU-load condition: key cells only
                        cell_budget = (
                            budget if args.quick else adaptive_budget(
                                test_file, file_size, block, pattern, nocache,
                                threads, floor=budget)
                        )
                        runs = [
                            measure(test_file, file_size, block, pattern, nocache,
                                    threads, cell_budget, seed=1000 + r)
                            for r in range(args.repeats)
                        ]
                        mbps = [r["mb_per_s"] for r in runs]
                        cell = {
                            "block_bytes": block,
                            "pattern": pattern,
                            "nocache": nocache,
                            "threads": threads,
                            "gpu_load": gpu_loaded,
                            "repeats": args.repeats,
                            "budget_bytes": cell_budget,
                            "mb_per_s_mean": statistics.mean(mbps),
                            "mb_per_s_median": statistics.median(mbps),
                            "mb_per_s_std": statistics.stdev(mbps) if len(mbps) > 1 else 0.0,
                            "iops_mean": statistics.mean(r["iops"] for r in runs),
                        }
                        results.append(cell)
                        print(
                            f"{'GPU+' if gpu_loaded else '    '}"
                            f"{pattern:>10} {block >> 10:>5} KiB "
                            f"nocache={int(nocache)} t={threads}: "
                            f"{cell['mb_per_s_mean']:9.1f} MB/s",
                            flush=True,
                        )

    run_cells(gpu_loaded=False)

    gpu = GpuLoad()
    gpu_iters = 0
    if gpu.available:
        # the warm-cache phase above populated the page cache; recreate the
        # file (uncached write) so the GPU-load cells are genuinely cold
        print("recreating test file to evict cached pages before GPU-load phase…", flush=True)
        create_test_file(test_file, file_size)
        print("re-running key cells under concurrent Metal (MLX) matmul load…", flush=True)
        with gpu:
            run_cells(gpu_loaded=True)
        gpu_iters = gpu.iterations
    else:
        print("MLX not available — skipping GPU-load condition", flush=True)

    iostat.__exit__()
    payload = {
        "experiment": "expH_ssd_feasibility",
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "manifest": manifest,
        "config": vars(args),
        "test_file_bytes": file_size,
        "gpu_load_matmul_iterations": gpu_iters,
        "thermal_level_after": collect_manifest()["thermal_level_at_collect"],
        "cells": results,
    }
    out_path = out_dir / "results.json"
    out_path.write_text(json.dumps(payload, indent=2))
    print(f"\nwrote {out_path}")


if __name__ == "__main__":
    main()