#!/usr/bin/env python3

"""Score the runner with a fixed ~20-second micro-suite (advisory only).

Three probes — CPU (sha256 over deterministic bytes), IO (write+fsync+read a

temp file), spawn (subprocess startup) — stamped with the runner identity from

oakbench.runner.  No thresholds and no exit-code judgment: calibration numbers

contextualize latency rows, they never gate them.

"""

from __future__ import annotations

import argparse

import hashlib

import json

import os

import random

import statistics

import subprocess

import sys

import tempfile

import time

from pathlib import Path

from typing import Any

ROOT = Path(__file__).resolve().parents[1]

sys.path.insert(0, str(ROOT / "scripts"))

from oakbench.runner import machine_profile, runner_class, runner_id  # noqa: E402

PROBE_BYTES = 64 * 1024 * 1024

QUICK_DIVISOR = 8

CPU_PASSES = 4

SPAWN_COUNT = 50

SPAWN_TIMEOUT_S = 30

DATA_SEED = 20260612

def parse_args() -> argparse.Namespace:

    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument("--path", type=Path, default=Path.cwd(), help="Directory for the IO probe temp file")

    parser.add_argument("--json", action="store_true", default=True, help="Emit JSON (always on; kept for symmetry)")

    parser.add_argument("--quick", action="store_true", help="Scale probe sizes down 8x for tests")

    return parser.parse_args()

def deterministic_bytes(size: int, seed: int = DATA_SEED) -> bytes:

    return random.Random(seed).randbytes(size)

def probe_cpu_sha256(data: bytes, passes: int = CPU_PASSES) -> float | None:

    """sha256 throughput in MiB/s over fixed in-process bytes."""

    if not data or passes < 1:

        return None

    start = time.perf_counter()

    for _ in range(passes):

        hashlib.sha256(data).digest()

    elapsed = time.perf_counter() - start

    if elapsed <= 0:

        return None

    return (len(data) / (1024 * 1024)) * passes / elapsed

def probe_io_write_read(directory: Path, data: bytes) -> float | None:

    """write+fsync+read throughput in MiB/s (total bytes moved); None if the path is unusable."""

    if not data:

        return None

    temp_path: Path | None = None

    try:

        start = time.perf_counter()

        with tempfile.NamedTemporaryFile(dir=directory, prefix="oakbench-calib-", delete=False) as handle:

            temp_path = Path(handle.name)

            handle.write(data)

            handle.flush()

            os.fsync(handle.fileno())

        read_back = temp_path.read_bytes()

        elapsed = time.perf_counter() - start

    except Exception:

        return None

    finally:

        try:

            if temp_path is not None:

                temp_path.unlink()

        except Exception:

            pass

    if elapsed <= 0 or len(read_back) != len(data):

        return None

    return (len(data) * 2 / (1024 * 1024)) / elapsed

def _spawn_once_ms() -> float | None:

    try:

        start = time.perf_counter()

        subprocess.run(

            [sys.executable, "-c", "pass"],

            stdout=subprocess.DEVNULL,

            stderr=subprocess.DEVNULL,

            timeout=SPAWN_TIMEOUT_S,

            check=True,

        )

        return (time.perf_counter() - start) * 1000.0

    except Exception:

        return None

def probe_spawn_median_ms(count: int = SPAWN_COUNT) -> float | None:

    samples = [sample for _ in range(count) if (sample := _spawn_once_ms()) is not None]

    return statistics.median(samples) if samples else None

def calibration_row(path: Path, quick: bool) -> dict[str, Any]:

    divisor = QUICK_DIVISOR if quick else 1

    started = time.perf_counter()

    profile = machine_profile(path)

    try:

        klass = runner_class(profile)

    except Exception:

        klass = None

    data = deterministic_bytes(PROBE_BYTES // divisor)

    row: dict[str, Any] = {

        "schema_version": 1,

        "runner_id": runner_id(profile),

        "runner_class": klass,

        "cpu_sha256_mib_per_s": probe_cpu_sha256(data),

        "io_write_read_mib_per_s": probe_io_write_read(path, data),

        "spawn_ms_median": probe_spawn_median_ms(max(1, SPAWN_COUNT // divisor)),

    }

    row["elapsed_s"] = time.perf_counter() - started

    return row

def main() -> int:

    args = parse_args()

    row = calibration_row(args.path, args.quick)

    print(json.dumps(row, sort_keys=True))

    return 0

if __name__ == "__main__":

    raise SystemExit(main())