#!/usr/bin/env python3
"""Print the tinker:// path of a stage's final checkpoint.

    python3 scripts/last_checkpoint.py runs/sft [--kind state|sampler]

`make all` uses this to hand stage N's output to stage N+1 without you having
to copy a run id out of the scrollback. Reads the checkpoints.jsonl the cookbook
writes into each run directory, preferring the record flagged `final`.
"""

import argparse
import json
import os
import sys


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("run_dir")
    ap.add_argument("--kind", choices=("state", "sampler"), default="state",
                    help="state resumes training; sampler is for inference/export")
    args = ap.parse_args()

    path = os.path.join(args.run_dir, "checkpoints.jsonl")
    if not os.path.exists(path):
        sys.exit("No checkpoints.jsonl in %s -- has that stage run?" % args.run_dir)

    records = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    records.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    if not records:
        sys.exit("%s is empty" % path)

    key = "%s_path" % args.kind
    # The cookbook marks the last checkpoint with {"name": "final", ...}; there
    # is no boolean "final" field. Checking r.get("final") never matched, so
    # this silently degraded to "whatever was written last" -- which is usually
    # right and would have hidden the bug until a run saved after its final.
    finals = [r for r in records if r.get("name") == "final" and r.get(key)]
    chosen = (finals or [r for r in records if r.get(key)])[-1:]
    if not chosen:
        sys.exit("No %s checkpoint recorded in %s" % (args.kind, path))
    print(chosen[0][key])


if __name__ == "__main__":
    main()