"""Sample jtop (jetson-stats) Python API → per-sample CSV rows.

Unlike tegrastats which is a stdout stream, jtop exposes a Python API
that emits a polled state dictionary. We poll at a caller-supplied
cadence and convert the relevant fields to CSV columns aligned with the
tegrastats output where the two overlap.

Schema (CSV columns):
    timestamp_utc_iso, ram_used_mb, ram_total_mb, gpu_load_pct,
    gpu_freq_mhz, cpu_load_avg_pct, soc_temp_c, gpu_temp_c, power_mw,
    extras_json

Usage:
    python3 jtop_parser.py --out out.csv --interval 1.0
"""

from __future__ import annotations

import argparse
import csv
import json
import time
from datetime import datetime, timezone

UTC = timezone.utc
from pathlib import Path


CSV_COLUMNS = (
    "timestamp_utc_iso",
    "ram_used_mb",
    "ram_total_mb",
    "gpu_load_pct",
    "gpu_freq_mhz",
    "cpu_load_avg_pct",
    "soc_temp_c",
    "gpu_temp_c",
    "power_mw",
    "extras_json",
)


def state_to_row(state: object) -> dict[str, object]:
    """Convert one jtop polled-state object to a CSV row.

    `state` is whatever `jtop.jtop().stats` returns; on real Jetson runs it
    is a `JtopStats` dataclass-ish object exposing `ram`, `gpu`, `cpu`,
    `temperature`, `power`. We extract defensively because jetson-stats
    schema has shifted across versions.
    """

    def _get(obj: object, *path: str, default: object = "") -> object:
        cur = obj
        for key in path:
            if cur is None:
                return default
            if isinstance(cur, dict):
                cur = cur.get(key, default)
            else:
                cur = getattr(cur, key, default)
        return cur if cur is not None else default

    row: dict[str, object] = {
        "timestamp_utc_iso": datetime.now(UTC).isoformat(timespec="milliseconds"),
        "ram_used_mb": _get(state, "ram", "used"),
        "ram_total_mb": _get(state, "ram", "tot"),
        "gpu_load_pct": _get(state, "gpu", "load"),
        "gpu_freq_mhz": _get(state, "gpu", "freq", "cur"),
        "cpu_load_avg_pct": _get(state, "cpu", "load_avg", default=""),
        "soc_temp_c": _get(state, "temperature", "SOC", default=""),
        "gpu_temp_c": _get(state, "temperature", "GPU", default=""),
        "power_mw": _get(state, "power", "total", default=""),
        "extras_json": "",
    }
    return row


def run(out_path: Path, interval_s: float, samples_max: int | None = None) -> int:
    """Poll jtop and write rows to ``out_path``. Returns rows written.

    On hosts without jetson-stats installed (e.g., unit-test runs on dev
    workstations), the function ImportError → emits a single "stub" row
    pointing at the missing dependency and exits. This keeps Tier-2 dry
    runs and CI smoke happy without forcing CI to install jetson-stats.
    """
    out_path.parent.mkdir(parents=True, exist_ok=True)
    rows_written = 0
    try:
        from jtop import jtop  # type: ignore[import-untyped]
    except ImportError as exc:
        with out_path.open("w", newline="", encoding="utf-8") as fh:
            writer = csv.DictWriter(fh, fieldnames=list(CSV_COLUMNS))
            writer.writeheader()
            writer.writerow(
                {
                    **{col: "" for col in CSV_COLUMNS},
                    "timestamp_utc_iso": datetime.now(UTC).isoformat(timespec="milliseconds"),
                    "extras_json": json.dumps({"stub": True, "missing_dep": "jetson-stats", "import_error": str(exc)}),
                }
            )
        return 1

    with jtop() as poll, out_path.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(CSV_COLUMNS))
        writer.writeheader()
        while poll.ok():
            row = state_to_row(poll.stats)
            writer.writerow(row)
            fh.flush()
            rows_written += 1
            if samples_max is not None and rows_written >= samples_max:
                break
            time.sleep(interval_s)
    return rows_written


def main() -> int:
    parser = argparse.ArgumentParser(description="Sample jtop → CSV.")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--interval", type=float, default=1.0, help="Poll interval in seconds.")
    parser.add_argument("--samples-max", type=int, default=None)
    args = parser.parse_args()
    n = run(args.out, args.interval, args.samples_max)
    print(f"jtop_parser: wrote {n} rows to {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())