"""Append one machine-readable record per fine-tuning run or evaluation to the lab notebook.

Purpose: Part 13's self-contained copy of the run-log format defined in Part 11, so that
         this part's scripts record a run 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, trl, peft 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 sftlog; sftlog.record(labbook="labbook.md", lab="part-13/train-lora", ...)
       or called from a shell script with the run's own fields as JSON on stdin:
           python3 sftlog.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 Part 11 defines. A record missing any of them is refused, because 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",
)

PACKAGES = ("torch", "transformers", "trl", "peft", "datasets", "bitsandbytes", "mlx", "mlx-lm")


def new_run_id() -> str:
    """Sorts by time and does not collide between two runs started in the same second."""
    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 is tied to the exact bytes it read.

    This is the field people leave out and then need: a month later the question is not
    which learning rate was used, which is in the script, but whether this run was before
    or after the dataset was fixed.
    """
    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, with -dirty when the tree had uncommitted edits."""
    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 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)
        return info
    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, ...] = PACKAGES) -> 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. Separate from writing so a caller can inspect it 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. A missing required field is an error, not a warning."""
    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 its id."""
    rec = build_record(**kwargs)
    append(rec, labbook)
    return rec


def main() -> None:
    parser = argparse.ArgumentParser(description="Run-log helper for Part 13's fine-tuning labs.")
    parser.add_argument("--record", action="store_true",
                        help="read this run's own fields as a JSON object 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()