"""Small helpers shared by benchmark entry points."""

from __future__ import annotations

import argparse
import json
import platform
import statistics
import time
from collections.abc import Callable
from typing import Any

import torch


def common_parser(description: str) -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--seed", type=int, default=1234)
    return parser


def synchronize(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def measure(
    operation: Callable[[], Any], *, iterations: int, warmup: int, device: torch.device
) -> dict[str, float]:
    for _ in range(warmup):
        operation()
    synchronize(device)
    durations = []
    for _ in range(iterations):
        start = time.perf_counter()
        operation()
        synchronize(device)
        durations.append(time.perf_counter() - start)
    return {
        "mean_ms": statistics.mean(durations) * 1000,
        "median_ms": statistics.median(durations) * 1000,
        "min_ms": min(durations) * 1000,
        "max_ms": max(durations) * 1000,
        "stdev_ms": statistics.stdev(durations) * 1000 if len(durations) > 1 else 0.0,
    }


def metadata(device: torch.device) -> dict[str, Any]:
    result: dict[str, Any] = {
        "python": platform.python_version(),
        "torch": torch.__version__,
        "device": str(device),
        "cuda_available": torch.cuda.is_available(),
    }
    if device.type == "cuda" and torch.cuda.is_available():
        result["gpu"] = torch.cuda.get_device_name(device)
    try:
        import dgl

        result["dgl"] = dgl.__version__
    except ImportError:
        result["dgl"] = None
    return result


def report(name: str, values: dict[str, Any]) -> None:
    print(json.dumps({"benchmark": name, **values}, sort_keys=True))