#!/usr/bin/env python3
"""Timestamped NVML sampling with power-rate and energy kept distinct."""

from __future__ import annotations

import argparse
import json
import time

import pynvml


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--seconds", type=float, default=5.0)
    parser.add_argument("--interval", type=float, default=0.1)
    args = parser.parse_args()

    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    start = time.monotonic()
    previous_t = start
    previous_watts = pynvml.nvmlDeviceGetPowerUsage(handle) / 1_000.0
    integrated_joules = 0.0
    samples: list[dict[str, float]] = []

    try:
        start_energy_mj = pynvml.nvmlDeviceGetTotalEnergyConsumption(handle)
    except pynvml.NVMLError:
        start_energy_mj = None

    while True:
        time.sleep(args.interval)
        now = time.monotonic()
        watts = pynvml.nvmlDeviceGetPowerUsage(handle) / 1_000.0
        integrated_joules += 0.5 * (previous_watts + watts) * (now - previous_t)
        samples.append({"t_s": now - start, "power_W": watts})
        previous_t, previous_watts = now, watts
        if now - start >= args.seconds:
            break

    try:
        end_energy_mj = pynvml.nvmlDeviceGetTotalEnergyConsumption(handle)
    except pynvml.NVMLError:
        end_energy_mj = None

    pynvml.nvmlShutdown()
    print(
        json.dumps(
            {
                "boundary": "gpu_device_not_rack_or_facility",
                "sample_interval_requested_s": args.interval,
                "integrated_sampled_energy_J": integrated_joules,
                "device_counter_energy_J": (
                    (end_energy_mj - start_energy_mj) / 1_000.0
                    if start_energy_mj is not None and end_energy_mj is not None
                    else None
                ),
                "samples": samples,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()