"""Append one machine-readable run record per training run to the lab notebook.

Purpose: the course's run log. Every training run from Part 11 onwards appends one
         JSON line to labbook.md describing what was run, on what, with what
         settings, and what came out, so that a result read a week later still
         means something.
Platform: all (pure Python; torch, transformers, trl, peft and mlx are only
          inspected for their version strings if they happen to be installed)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer; the lab notebook from Part 1 exists (it is created
         if it does not); git is optional and only used to record the commit of
         the configuration file.

Usage: imported by the training scripts, or run directly from a shell script:
       python runlog.py --demo --labbook labbook.md
       python runlog.py --record --labbook labbook.md < fields.json

The --record form reads a JSON object with the keys lab, model, dataset,
hyperparameters, seed, losses, scores and notes, and fills in the run id, date,
configuration commit, hardware and package versions itself.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import secrets
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# Every field below is required in a run record. A run that cannot fill one in
# writes null rather than leaving the key out, so that a reader always knows the
# difference between "not recorded" and "not applicable".
FIELDS = (
    "run_id",
    "lab",
    "date",
    "config_commit",
    "model",
    "dataset",
    "hyperparameters",
    "seed",
    "hardware",
    "versions",
    "losses",
    "scores",
    "notes",
)


def new_run_id() -> str:
    """A short identifier that sorts by time and does not collide between runs."""
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    return f"{stamp}-{secrets.token_hex(3)}"


def file_sha256(path: str | os.PathLike[str], chunk: int = 1 << 20) -> str | None:
    """Content hash of a dataset file, so a run can be tied to the exact bytes it read."""
    p = Path(path)
    if not p.is_file():
        return None
    digest = hashlib.sha256()
    with p.open("rb") as handle:
        while True:
            block = handle.read(chunk)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def git_commit(path: str | os.PathLike[str] = ".") -> str | None:
    """The commit the configuration is at, or None outside a repository.

    Records whether the tree was dirty, because "commit abc1234" is misleading
    if the file on disk had uncommitted edits when the run started.
    """
    if shutil.which("git") is None:
        return None
    target = Path(path)
    cwd = target if target.is_dir() else target.parent
    try:
        rev = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=cwd, capture_output=True, text=True, check=True, timeout=10,
        ).stdout.strip()
        dirty = subprocess.run(
            ["git", "status", "--porcelain"],
            cwd=cwd, capture_output=True, text=True, check=True, timeout=10,
        ).stdout.strip()
    except (subprocess.SubprocessError, OSError):
        return None
    return f"{rev}-dirty" if dirty else rev


def describe_hardware() -> dict[str, Any]:
    """What the run was executed on, as far as it can be established without extra packages."""
    info: dict[str, Any] = {
        "os": f"{platform.system()} {platform.release()}",
        "machine": platform.machine(),
        "python": platform.python_version(),
        "accelerator": "cpu",
        "device_name": None,
    }
    try:
        import torch  # noqa: PLC0415 - optional, and only for reporting
    except ImportError:
        return info
    if torch.cuda.is_available():
        info["accelerator"] = "cuda"
        info["device_name"] = torch.cuda.get_device_name(0)
    else:
        mps = getattr(torch.backends, "mps", None)
        if mps is not None and mps.is_available():
            info["accelerator"] = "mps"
            info["device_name"] = platform.processor() or "Apple silicon"
    return info


def package_versions(names: tuple[str, ...] = ("torch", "transformers", "trl", "peft", "datasets", "mlx")) -> dict[str, str | None]:
    """Version strings for the packages that decide what a run actually did."""
    from importlib.metadata import PackageNotFoundError, version  # noqa: PLC0415

    out: dict[str, str | None] = {}
    for name in names:
        try:
            out[name] = version(name)
        except PackageNotFoundError:
            out[name] = None
    return out


def build_record(
    lab: str,
    model: str,
    dataset: dict[str, Any],
    hyperparameters: dict[str, Any],
    seed: int,
    losses: dict[str, Any] | None = None,
    scores: dict[str, Any] | None = None,
    config_path: str | os.PathLike[str] | None = None,
    notes: str | None = None,
) -> dict[str, Any]:
    """Assemble the record. Kept separate from writing so it can be inspected first."""
    return {
        "run_id": new_run_id(),
        "lab": lab,
        "date": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "config_commit": git_commit(config_path if config_path is not None else "."),
        "model": model,
        "dataset": dataset,
        "hyperparameters": hyperparameters,
        "seed": seed,
        "hardware": describe_hardware(),
        "versions": package_versions(),
        "losses": losses or {},
        "scores": scores or {},
        "notes": notes,
    }


def append(record: dict[str, Any], labbook: str | os.PathLike[str] = "labbook.md") -> Path:
    """Append one JSON line. Missing keys are an error: a partial record is worse than none."""
    missing = [f for f in FIELDS if f not in record]
    if missing:
        raise ValueError(f"run record is missing required fields: {', '.join(missing)}")
    path = Path(labbook)
    if not path.exists():
        path.write_text("# Lab notebook\n\n## Results\n\n", encoding="utf-8")
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")
    return path


def record(labbook: str | os.PathLike[str] = "labbook.md", **kwargs: Any) -> dict[str, Any]:
    """Build and append in one call; returns the record so a caller can print it."""
    rec = build_record(**kwargs)
    append(rec, labbook)
    return rec


def main() -> None:
    parser = argparse.ArgumentParser(description="Run-log helper for the training labs.")
    parser.add_argument("--demo", action="store_true", help="append an example record and print it")
    parser.add_argument("--record", action="store_true", help="read the run's own fields as JSON on stdin")
    parser.add_argument("--labbook", default="labbook.md")
    args = parser.parse_args()
    if args.record:
        fields = json.load(sys.stdin)
        allowed = {"lab", "model", "dataset", "hyperparameters", "seed", "losses", "scores", "config_path", "notes"}
        unknown = set(fields) - allowed
        if unknown:
            raise SystemExit(f"unknown field(s) on stdin: {', '.join(sorted(unknown))}")
        rec = record(labbook=args.labbook, **fields)
        print(f"recorded run {rec['run_id']} in {args.labbook}")
        return
    if not args.demo:
        print(__doc__)
        print("Fields in every record:", ", ".join(FIELDS))
        return
    rec = record(
        labbook=args.labbook,
        lab="part-11/runlog-demo",
        model="none",
        dataset={"path": None, "sha256": None, "examples": 0},
        hyperparameters={"note": "example record written by runlog.py --demo"},
        seed=0,
        losses={},
        scores={},
        notes="delete this line from labbook.md once you have seen it",
    )
    json.dump(rec, sys.stdout, indent=2, sort_keys=True)
    print()


if __name__ == "__main__":
    main()