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

Purpose: Part 17's self-contained copy of the run-log format defined in Part 11, so that
         this part's scripts record a data-preparation run, a draft-training run or a
         serving measurement identically without depending on Part 11's files being on
         the path. Every field is filled in or written as null, because a reader a month
         later has to be able to tell "not recorded" from "not applicable".
Platform: all (standard library only; torch, transformers and mlx are inspected for their
          version strings only if they happen to be installed)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer. The lab notebook is created if it does not exist. git is
         optional and is used only to record the commit the configuration was at.

Usage: imported by this part's Python scripts:
           import draftlog; draftlog.record(labbook="labbook.md", lab="part-17/make-draft-data", ...)
       or called from a shell script with the run's own fields as JSON on stdin:
           python3 draftlog.py --record --labbook labbook.md < fields.json
       or run with no arguments to print the field list and exit.
"""
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

# The same field set Parts 11 and 13 define, so a Part 17 line can be read by the same
# tooling. A record missing any of them is refused: a partial record is harder to
# interpret than no record at all.
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 data 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."""
    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", "accelerate", "mlx", "mlx-lm"),
) -> 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 Part 17 draft-model lab.")
    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 not args.record:
        print(__doc__)
        print("Fields in every record:", ", ".join(FIELDS))
        return
    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}")


if __name__ == "__main__":
    main()