"""Append one machine-readable record per distillation stage to the lab notebook.

Purpose: Part 15's self-contained copy of the run-log format defined in Part 11 and
         reused in Part 13, extended with the two fields distillation adds: the
         teacher a stage used and what the stage cost in tokens, seconds and
         watt-hours. Every stage of the pipeline appends one line, so the whole
         run reads back as an ordered story rather than a directory of artefacts.
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. nvidia-smi or rocm-smi are optional and are used only by
         sample_power(), which returns None when neither is present.

Usage: imported by this part's Python scripts:
           import distillog
           distillog.record(labbook="labbook.md", lab="part-15/generate", ...)
       or called from a shell script with the stage's own fields as JSON on stdin:
           python3 distillog.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 re
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 record. A stage that cannot fill one in
# writes null rather than omitting the key, so a reader always knows the
# difference between "not recorded" and "not applicable".
FIELDS = (
    "run_id",
    "lab",
    "stage",
    "date",
    "config_commit",
    "teacher",
    "student",
    "dataset",
    "hyperparameters",
    "seed",
    "hardware",
    "versions",
    "cost",
    "losses",
    "scores",
    "notes",
)

# The cost block is the reason this file exists rather than Part 13's sftlog.py.
# A distillation result that does not say what it cost cannot be compared with
# the alternative that would have cost less.
COST_FIELDS = ("prompt_tokens", "completion_tokens", "seconds", "watt_hours")


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 stage 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 config_sha256(config: dict[str, Any]) -> str:
    """Hash of a configuration object, sorted so that key order cannot change it.

    Two runs of the pipeline with the same configuration hash read the same
    settings; two with different hashes did not, whatever the file names say.
    """
    return hashlib.sha256(json.dumps(config, sort_keys=True).encode("utf-8")).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 stage 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 stage ran 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", "distilabel"),
) -> dict[str, str | None]:
    """Version strings for the packages that decide what a stage 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 sample_power() -> float | None:
    """Instantaneous accelerator power draw in watts, or None where it is not reported.

    Track S and Track N expose it through nvidia-smi; Track X exposes it through
    rocm-smi. Apple silicon does not report a comparable per-accelerator figure
    to an unprivileged process, so Track M records None and the page says so.
    One sample is not an energy measurement: generate-teacher-data.py averages
    samples over the run and multiplies by the elapsed hours.
    """
    if shutil.which("nvidia-smi"):
        try:
            out = subprocess.run(
                ["nvidia-smi", "--query-gpu=power.draw", "--format=csv,noheader,nounits"],
                capture_output=True, text=True, check=True, timeout=10,
            ).stdout.strip().splitlines()
            values = [float(v) for v in out if re.fullmatch(r"\s*[0-9.]+\s*", v)]
            if values:
                return round(sum(values), 1)
        except (subprocess.SubprocessError, OSError, ValueError):
            return None
    if shutil.which("rocm-smi"):
        try:
            out = subprocess.run(
                ["rocm-smi", "--showpower", "--json"],
                capture_output=True, text=True, check=True, timeout=10,
            ).stdout
            data = json.loads(out)
            values = []
            for card in data.values():
                for key, value in card.items():
                    if "power" in key.lower():
                        try:
                            values.append(float(str(value).split()[0]))
                        except (TypeError, ValueError):
                            continue
            if values:
                return round(sum(values), 1)
        except (subprocess.SubprocessError, OSError, ValueError, json.JSONDecodeError):
            return None
    return None


def build_cost(
    prompt_tokens: int | None = None,
    completion_tokens: int | None = None,
    seconds: float | None = None,
    mean_watts: float | None = None,
) -> dict[str, Any]:
    """The four cost numbers, with watt-hours derived rather than typed.

    mean_watts is the average of sample_power() readings taken during the stage.
    Where no reading was available the energy figure is null, which is the honest
    answer: a missing measurement is not zero.
    """
    watt_hours = None
    if mean_watts is not None and seconds is not None:
        watt_hours = round(mean_watts * (seconds / 3600.0), 2)
    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "seconds": round(seconds, 1) if seconds is not None else None,
        "watt_hours": watt_hours,
    }


def build_record(
    lab: str,
    stage: str,
    teacher: dict[str, Any] | None = None,
    student: dict[str, Any] | None = None,
    dataset: dict[str, Any] | None = None,
    hyperparameters: dict[str, Any] | None = None,
    seed: int | None = None,
    cost: dict[str, Any] | None = None,
    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."""
    filled_cost = dict(cost or {})
    for key in COST_FIELDS:
        filled_cost.setdefault(key, None)
    return {
        "run_id": new_run_id(),
        "lab": lab,
        "stage": stage,
        "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 "."),
        "teacher": teacher,
        "student": student,
        "dataset": dataset or {},
        "hyperparameters": hyperparameters or {},
        "seed": seed,
        "hardware": describe_hardware(),
        "versions": package_versions(),
        "cost": filled_cost,
        "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 read_stages(labbook: str | os.PathLike[str], lab_prefix: str = "part-15/") -> list[dict[str, Any]]:
    """Every Part 15 record in the notebook, oldest first.

    The project's pipeline runner uses this to print what has already been done
    and what a resumed run still has to do.
    """
    path = Path(labbook)
    if not path.is_file():
        return []
    out = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line.startswith("{"):
            continue
        try:
            rec = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(rec, dict) and str(rec.get("lab", "")).startswith(lab_prefix):
            out.append(rec)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description="Run-log helper for the distillation labs.")
    parser.add_argument("--record", action="store_true", help="read the stage's own fields as JSON on stdin")
    parser.add_argument("--stages", action="store_true", help="list the Part 15 stages already in the notebook")
    parser.add_argument("--power", action="store_true", help="print one accelerator power reading and exit")
    parser.add_argument("--labbook", default="labbook.md")
    args = parser.parse_args()

    if args.power:
        watts = sample_power()
        print("not reported on this machine" if watts is None else f"{watts} W")
        return
    if args.stages:
        for rec in read_stages(args.labbook):
            cost = rec.get("cost") or {}
            print(f"{rec['date']}  {rec['stage']:<10}  {rec['run_id']}  "
                  f"{cost.get('completion_tokens') or '-'} completion tokens  "
                  f"{cost.get('seconds') or '-'} s")
        return
    if args.record:
        fields = json.load(sys.stdin)
        allowed = {"lab", "stage", "teacher", "student", "dataset", "hyperparameters",
                   "seed", "cost", "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 {rec['stage']} as {rec['run_id']} in {args.labbook}")
        return

    print(__doc__)
    print("Fields in every record:", ", ".join(FIELDS))
    print("Fields in every cost block:", ", ".join(COST_FIELDS))


if __name__ == "__main__":
    main()