#!/usr/bin/env python3

"""Cross-check harness timings with optional external `hyperfine`.

The benchmark harness stays stdlib-only. `hyperfine` is optional: when absent,

callers get either a clear diagnostic or a structured returncode-77 skip row.

"""

from __future__ import annotations

import argparse

import json

import os

import shlex

import shutil

import subprocess

import sys

import tempfile

from pathlib import Path

from typing import Any, Callable, Optional, Sequence

from oakbench.environment import host_home_environment

SKIP_RETURNCODE = 77

MEASUREMENT_SOURCE = "external_hyperfine"

def hyperfine_path(executable: str = "hyperfine") -> Optional[str]:

    return shutil.which(executable)

def skip_row(

    *,

    command: Sequence[str],

    reason: str,

    bench_id: str = "manual",

    scenario: str = "external-timing",

    operation: str = "hyperfine.crosscheck",

    run: int = 1,

    subject: str = "unknown",

    subject_kind: str = "unknown",

) -> dict[str, Any]:

    return {

        "schema_version": 1,

        "bench_id": bench_id,

        "profile": "crosscheck",

        "scenario": scenario,

        "operation": operation,

        "run": run,

        "subject": subject,

        "subject_kind": subject_kind,

        "elapsed_ms": 0.0,

        "returncode": SKIP_RETURNCODE,

        "command": list(command),

        "skipped": True,

        "skip_reason": reason,

        "measurement_source": MEASUREMENT_SOURCE,

    }

def hyperfine_command(executable: str, command: Sequence[str], export_json: Path, *, runs: int, warmup: int) -> list[str]:

    return [

        executable,

        "--export-json",

        str(export_json),

        "--runs",

        str(runs),

        "--warmup",

        str(warmup),

        shlex.join(command),

    ]

def parse_hyperfine_json(payload: str) -> dict[str, Any]:

    data = json.loads(payload)

    results = data.get("results")

    if not isinstance(results, list) or not results:

        raise ValueError("hyperfine JSON did not contain any results")

    first = results[0]

    if not isinstance(first, dict):

        raise ValueError("hyperfine JSON result is not an object")

    return first

def measurement_row(

    *,

    command: Sequence[str],

    hyperfine_result: dict[str, Any],

    bench_id: str = "manual",

    scenario: str = "external-timing",

    operation: str = "hyperfine.crosscheck",

    run: int = 1,

    subject: str = "unknown",

    subject_kind: str = "unknown",

) -> dict[str, Any]:

    mean = _number_or_none(hyperfine_result.get("mean"))

    median = _number_or_none(hyperfine_result.get("median"))

    stddev = _number_or_none(hyperfine_result.get("stddev"))

    minimum = _number_or_none(hyperfine_result.get("min"))

    maximum = _number_or_none(hyperfine_result.get("max"))

    times = hyperfine_result.get("times")

    times_ms = None

    if isinstance(times, list):

        times_ms = [round(float(value) * 1000.0, 6) for value in times if _number_or_none(value) is not None]

    elapsed_seconds = median if median is not None else mean

    return {

        "schema_version": 1,

        "bench_id": bench_id,

        "profile": "crosscheck",

        "scenario": scenario,

        "operation": operation,

        "run": run,

        "subject": subject,

        "subject_kind": subject_kind,

        "elapsed_ms": round((elapsed_seconds or 0.0) * 1000.0, 6),

        "returncode": 0,

        "command": list(command),

        "measurement_source": MEASUREMENT_SOURCE,

        "hyperfine_command": hyperfine_result.get("command"),

        "hyperfine_mean_ms": _seconds_to_ms(mean),

        "hyperfine_median_ms": _seconds_to_ms(median),

        "hyperfine_stddev_ms": _seconds_to_ms(stddev),

        "hyperfine_min_ms": _seconds_to_ms(minimum),

        "hyperfine_max_ms": _seconds_to_ms(maximum),

        "hyperfine_times_ms": times_ms,

    }

def run_hyperfine(

    command: Sequence[str],

    *,

    cwd: Path,

    runs: int,

    warmup: int,

    executable: str = "hyperfine",

    env: Optional[dict[str, str]] = None,

    runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,

) -> dict[str, Any]:

    resolved = hyperfine_path(executable)

    if resolved is None:

        raise HyperfineUnavailable(

            f"hyperfine not found on PATH; install hyperfine or emit a returncode-{SKIP_RETURNCODE} skip row"

        )

    with tempfile.TemporaryDirectory(prefix="oakbench-hyperfine-") as tmp:

        export_json = Path(tmp) / "hyperfine.json"

        proc = runner(

            hyperfine_command(resolved, command, export_json, runs=runs, warmup=warmup),

            cwd=cwd,

            env=host_home_environment(env),

            stdout=subprocess.PIPE,

            stderr=subprocess.PIPE,

            text=True,

        )

        if proc.returncode != 0:

            raise HyperfineFailed(proc.returncode, proc.stderr.strip() or proc.stdout.strip())

        return parse_hyperfine_json(export_json.read_text())

class HyperfineUnavailable(RuntimeError):

    pass

class HyperfineFailed(RuntimeError):

    def __init__(self, returncode: int, message: str) -> None:

        super().__init__(message)

        self.returncode = returncode

def _number_or_none(value: Any) -> Optional[float]:

    if value is None:

        return None

    try:

        return float(value)

    except (TypeError, ValueError):

        return None

def _seconds_to_ms(value: Optional[float]) -> Optional[float]:

    return None if value is None else round(value * 1000.0, 6)

def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:

    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument("--cwd", type=Path, default=Path.cwd())

    parser.add_argument("--runs", type=int, default=10)

    parser.add_argument("--warmup", type=int, default=1)

    parser.add_argument("--hyperfine", default="hyperfine")

    parser.add_argument("--jsonl", action="store_true", help="print one JSON row, including skip rows")

    parser.add_argument("--bench-id", default="manual")

    parser.add_argument("--scenario", default="external-timing")

    parser.add_argument("--operation", default="hyperfine.crosscheck")

    parser.add_argument("--run", type=int, default=1)

    parser.add_argument("--subject", default="unknown")

    parser.add_argument("--subject-kind", default="unknown")

    parser.add_argument("command", nargs=argparse.REMAINDER, help="command to measure, usually after --")

    args = parser.parse_args(argv)

    if args.command and args.command[0] == "--":

        args.command = args.command[1:]

    if not args.command:

        parser.error("command is required; pass it after --")

    if args.runs <= 0:

        parser.error("--runs must be positive")

    if args.warmup < 0:

        parser.error("--warmup must be non-negative")

    return args

def main(argv: Optional[Sequence[str]] = None) -> int:

    args = parse_args(argv)

    row_kwargs = {

        "bench_id": args.bench_id,

        "scenario": args.scenario,

        "operation": args.operation,

        "run": args.run,

        "subject": args.subject,

        "subject_kind": args.subject_kind,

    }

    try:

        result = run_hyperfine(

            args.command,

            cwd=args.cwd,

            runs=args.runs,

            warmup=args.warmup,

            executable=args.hyperfine,

            env=host_home_environment(),

        )

    except HyperfineUnavailable as exc:

        if args.jsonl:

            print(json.dumps(skip_row(command=args.command, reason=str(exc), **row_kwargs), sort_keys=True))

        else:

            print(str(exc), file=sys.stderr)

        return SKIP_RETURNCODE

    except HyperfineFailed as exc:

        print(f"hyperfine failed with exit code {exc.returncode}: {exc}", file=sys.stderr)

        return exc.returncode or 1

    print(json.dumps(measurement_row(command=args.command, hyperfine_result=result, **row_kwargs), sort_keys=True))

    return 0

if __name__ == "__main__":

    raise SystemExit(main())