"""
Lab-notebook helper. Everything the run learns goes to logs/ as markdown +
figures, so the record survives the process that produced it.

Layout:
  logs/lab_notebook.md          chronological running journal (append-only)
  logs/experiments/<name>.md    one long report per experiment
  logs/figures/<name>.png       figures
  logs/disk.md                  space accounting
"""

from __future__ import annotations

import json
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
LOGS = ROOT / "logs"
FIGS = LOGS / "figures"
EXPS = LOGS / "experiments"
NOTEBOOK = LOGS / "lab_notebook.md"


def _ts() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%SZ")


def _ensure() -> None:
    for d in (LOGS, FIGS, EXPS):
        d.mkdir(parents=True, exist_ok=True)


def note(title: str, body: str = "", level: str = "INFO") -> None:
    """Append a timestamped entry to the running journal."""
    _ensure()
    with open(NOTEBOOK, "a") as f:
        f.write(f"\n### [{_ts()}] {level} — {title}\n\n")
        if body:
            f.write(body.rstrip() + "\n")


def table(rows: list[dict], cols: list[str] | None = None) -> str:
    """Render list-of-dicts as a markdown table."""
    if not rows:
        return "_(empty)_\n"
    cols = cols or list(rows[0].keys())

    def fmt(v):
        if isinstance(v, float):
            return f"{v:.4g}"
        return str(v)

    out = ["| " + " | ".join(cols) + " |",
           "|" + "|".join("---" for _ in cols) + "|"]
    for r in rows:
        out.append("| " + " | ".join(fmt(r.get(c, "")) for c in cols) + " |")
    return "\n".join(out) + "\n"


def write_report(name: str, content: str) -> Path:
    """Write/overwrite one experiment's full report."""
    _ensure()
    p = EXPS / f"{name}.md"
    p.write_text(content)
    return p


def gpu_snapshot() -> dict:
    try:
        q = ("--query-gpu=memory.used,memory.total,utilization.gpu,temperature.gpu"
             ",power.draw")
        out = subprocess.run(
            ["nvidia-smi", q, "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=20).stdout.strip()
        u, t, g, temp, pw = [x.strip() for x in out.split(",")]
        return {"vram_used_mb": int(u), "vram_total_mb": int(t),
                "gpu_util_pct": int(g), "temp_c": int(temp), "power_w": float(pw)}
    except Exception as e:
        return {"error": str(e)}


def disk_snapshot(paths: list[str] | None = None) -> dict:
    paths = paths or [str(ROOT)]
    out = {}
    try:
        df = subprocess.run(["df", "-BM", "--output=avail,used,size", str(ROOT)],
                            capture_output=True, text=True, timeout=20).stdout
        parts = df.strip().splitlines()[-1].split()
        out["fs_avail_mb"] = int(parts[0].rstrip("M"))
        out["fs_used_mb"] = int(parts[1].rstrip("M"))
    except Exception as e:
        out["df_error"] = str(e)
    for p in paths:
        try:
            du = subprocess.run(["du", "-sm", p], capture_output=True,
                                text=True, timeout=120).stdout.split()[0]
            out[f"du_mb:{Path(p).name}"] = int(du)
        except Exception:
            pass
    return out


def checkpoint(stage: str, extra: dict | None = None) -> dict:
    """One-line health snapshot appended to the journal."""
    snap = {"stage": stage, "gpu": gpu_snapshot(), "disk": disk_snapshot()}
    if extra:
        snap.update(extra)
    note(f"checkpoint: {stage}", "```json\n" + json.dumps(snap, indent=1) + "\n```")
    return snap


class Timer:
    def __init__(self, label: str):
        self.label = label

    def __enter__(self):
        self.t0 = time.time()
        return self

    def __exit__(self, *a):
        self.dt = time.time() - self.t0
        note(f"timing: {self.label}", f"`{self.dt:.1f}s` ({self.dt/60:.1f} min)")