"""Append one machine-readable record per collection, training or evaluation run to the notebook.

Purpose: Part 27's self-contained copy of the run-log format defined in Part 11, so that
    this part's scripts record a run identically without needing Part 11's or Part 13's
    files 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". The extra
    field this part adds is `data_lineage`: the trajectory file a training set came from,
    its hash, and the scrub and filter reports that stand between them.
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 agentlog
           agentlog.record(labbook="labbook.md", lab="part-27/train-agent-sft", ...)
       or called from a shell script with the run's own fields as JSON on standard input:
           python3 agentlog.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 subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# The same field set Part 11 defines, plus data_lineage. 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", "data_lineage",
    "hyperparameters", "seed", "hardware", "versions", "losses", "scores", "notes",
)

PACKAGES = ("torch", "transformers", "trl", "peft", "datasets", "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 file, so a run is tied to the exact bytes it read.

    A month later the question is rarely which learning rate was used, which is in the
    script. It is whether this run was before or after the scrub, and the hash answers it.
    """
    p = Path(path)
    if not p.is_file():
        return None
    digest = hashlib.sha256()
    with p.open("rb") as handle:
        for block in iter(lambda: handle.read(chunk), b""):
            digest.update(block)
    return digest.hexdigest()


def package_versions() -> dict[str, str | None]:
    """Version strings for the packages that decide what a training run actually did."""
    out: dict[str, str | None] = {}
    try:
        from importlib import metadata
    except ImportError:  # pragma: no cover - Python 3.7 and earlier only
        return {name: None for name in PACKAGES}
    for name in PACKAGES:
        try:
            out[name] = metadata.version(name)
        except Exception:
            out[name] = None
    return out


def git_commit(path: str | os.PathLike[str] | None) -> str | None:
    """The commit the configuration was at, when the configuration lives in git."""
    if path is None:
        return None
    directory = Path(path).resolve()
    directory = directory if directory.is_dir() else directory.parent
    try:
        done = subprocess.run(
            ["git", "-C", str(directory), "rev-parse", "--short", "HEAD"],
            capture_output=True, text=True, timeout=10, check=False,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    return done.stdout.strip() or None


def hardware() -> dict[str, Any]:
    """What the run happened on, at the level of detail that changes a conclusion."""
    info: dict[str, Any] = {
        "platform": platform.platform(),
        "machine": platform.machine(),
        "python": platform.python_version(),
        "accelerator": None,
        "accelerator_count": None,
    }
    try:
        import torch  # noqa: PLC0415 - optional and only inspected
    except Exception:
        return info
    try:
        if torch.cuda.is_available():
            info["accelerator"] = torch.cuda.get_device_name(0)
            info["accelerator_count"] = torch.cuda.device_count()
        elif getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
            info["accelerator"] = "Apple silicon (Metal, PyTorch MPS)"
            info["accelerator_count"] = 1
        else:
            info["accelerator"] = "cpu"
            info["accelerator_count"] = 0
    except Exception:
        pass
    return info


def lineage(trajectories: str | os.PathLike[str] | None = None,
            scrub_report: str | os.PathLike[str] | None = None,
            filter_report: str | os.PathLike[str] | None = None,
            **extra: Any) -> dict[str, Any]:
    """Where the training data came from, as paths plus content hashes.

    This is the field that makes a trajectory fine-tune auditable. Without it, a data set
    and a model are two files with no stated relationship, and the question "was this
    trained on scrubbed data" has no answer that can be checked.
    """
    out: dict[str, Any] = {
        "trajectories": str(trajectories) if trajectories else None,
        "trajectories_sha256": file_sha256(trajectories) if trajectories else None,
        "scrub_report": str(scrub_report) if scrub_report else None,
        "scrub_report_sha256": file_sha256(scrub_report) if scrub_report else None,
        "filter_report": str(filter_report) if filter_report else None,
        "filter_report_sha256": file_sha256(filter_report) if filter_report else None,
    }
    out.update(extra)
    return out


def record(labbook: str | os.PathLike[str], lab: str, **fields: Any) -> dict[str, Any]:
    """Build one record, append it to the notebook as a JSON line, and return it."""
    row: dict[str, Any] = {
        "run_id": fields.pop("run_id", None) or new_run_id(),
        "lab": lab,
        "date": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "config_commit": fields.pop("config_commit", None) or git_commit(fields.pop("config_path", None)),
        "model": fields.pop("model", None),
        "dataset": fields.pop("dataset", None),
        "data_lineage": fields.pop("data_lineage", None),
        "hyperparameters": fields.pop("hyperparameters", None),
        "seed": fields.pop("seed", None),
        "hardware": fields.pop("hardware", None) or hardware(),
        "versions": fields.pop("versions", None) or package_versions(),
        "losses": fields.pop("losses", None),
        "scores": fields.pop("scores", None),
        "notes": fields.pop("notes", None),
    }
    # Anything else the caller passed is kept rather than dropped: a script that records
    # one more number should not have to change this module.
    row.update(fields)
    missing = [name for name in FIELDS if name not in row]
    if missing:
        raise ValueError(f"run record is missing {missing}; every field is written, even as null")

    path = Path(labbook)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(row, default=str) + "\n")
    return row


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--record", action="store_true",
                        help="read a JSON object of fields from standard input and append it")
    parser.add_argument("--labbook", default="labbook.md")
    parser.add_argument("--lab", default="part-27/unnamed")
    args = parser.parse_args()

    if not args.record:
        print("Fields written on every record:")
        for name in FIELDS:
            print(f"  {name}")
        print("\nPass --record with a JSON object on standard input to append one.")
        return

    payload = json.load(sys.stdin)
    if not isinstance(payload, dict):
        sys.exit("standard input must be a JSON object of fields")
    written = record(labbook=args.labbook, lab=payload.pop("lab", args.lab), **payload)
    print(f"recorded run {written['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()