#!/usr/bin/env python3
"""Appends one line per model batch to the track's ledger on the shared NAS, so a landing
can say what the work behind it took (bench/feature-announce.py).

    bench/ledger.py [--print] <measure dir>

<measure dir> is the directory bench/batch-lane.sh batch_measure_start wrote: phases.tsv
(start and end time), run.stamped.log (the model-run binary's output, one unix time before
each line) and cpu-phases.txt (bench/cpu-phases.py's table; its "whole run" row gives the
lane CPU-seconds). batch_measure_stop runs this after writing cpu-phases.txt.

From the run log, per model: candidates drawn and draw seconds ("timing: <model>: drew N
candidate(s) in Xs"), compiled ("<model>: A/B candidate(s) compiled"), completion tokens
(the "(<model>): i/n received, K completion tokens" lines) and the best score ("<model>: N
scored, best <candidate> at P/R"). A run of several batches sums the counts and keeps the
highest score.

The line goes to <dir>/<track>.jsonl, where <dir> is NEV_LEDGER_DIR, else
$NEV_LANE_SHARED/ledger when that directory exists (every lane mounts the NAS there), and
<track> is NEV_LEDGER_TRACK or the name of the checkout's top directory (on a lane the
worktree copy keeps its name). With neither directory it prints one line saying so and
writes nothing. --print writes the record to stdout instead. It never fails a batch:
every error is one line on stderr and exit 0. The record holds no host names or addresses.

Record: {"format": 1, "track", "start", "end" (unix seconds), "cpu_seconds",
"models": {<model>: {"candidates", "compiled", "completion_tokens", "draw_seconds",
"best": {"passed", "run"} or null}}}.
"""

import json
import os
import re
import subprocess
import sys
from pathlib import Path

DREW = re.compile(r"^timing: (\S+): drew (\d+) candidate\(s\) in ([0-9.]+)s")
COMPILED = re.compile(r"^(\S+): (\d+)/(\d+) candidate\(s\) compiled$")
RECEIVED = re.compile(r"^\S+ \((\S+)\): \d+/\d+ received, (\d+) completion tokens")
BEST = re.compile(r"^(\S+): \d+ scored, best \S+ at (\d+)/(\d+)")
TRACK_NAME = re.compile(r"^[A-Za-z0-9._-]+$")


def strip_stamp(line):
    parts = line.split(" ", 1)
    try:
        float(parts[0])
    except ValueError:
        return line
    return parts[1] if len(parts) == 2 else ""


def parse_log(lines):
    models = {}

    def model(name):
        return models.setdefault(name, {
            "candidates": 0, "compiled": 0, "completion_tokens": 0,
            "draw_seconds": 0.0, "best": None,
        })

    for raw in lines:
        line = strip_stamp(raw.rstrip("\n"))
        m = DREW.match(line)
        if m:
            e = model(m.group(1))
            e["candidates"] += int(m.group(2))
            e["draw_seconds"] = round(e["draw_seconds"] + float(m.group(3)), 1)
            continue
        m = COMPILED.match(line)
        if m:
            model(m.group(1))["compiled"] += int(m.group(2))
            continue
        m = RECEIVED.match(line)
        if m:
            model(m.group(1))["completion_tokens"] += int(m.group(2))
            continue
        m = BEST.match(line)
        if m:
            e = model(m.group(1))
            passed, run = int(m.group(2)), int(m.group(3))
            if e["best"] is None or passed * e["best"]["run"] > e["best"]["passed"] * run:
                e["best"] = {"passed": passed, "run": run}
    return models


def span(phases_path):
    times = []
    with open(phases_path) as f:
        for line in f:
            try:
                times.append(float(line.split(" ", 1)[0]))
            except ValueError:
                pass
    return (round(min(times), 1), round(max(times), 1)) if times else (None, None)


def cpu_seconds(table_path):
    try:
        with open(table_path) as f:
            for line in f:
                if line.startswith("whole run"):
                    return round(float(line.split()[3]), 1)
    except (OSError, ValueError, IndexError):
        pass
    return None


def track_name(measure_dir):
    name = os.environ.get("NEV_LEDGER_TRACK", "")
    if not name:
        try:
            top = subprocess.run(["git", "-C", str(measure_dir), "rev-parse", "--show-toplevel"],
                                 capture_output=True, text=True, check=True).stdout.strip()
            name = Path(top).name
        except (OSError, subprocess.CalledProcessError):
            name = ""
    return name if TRACK_NAME.match(name) else ""


def ledger_dir():
    d = os.environ.get("NEV_LEDGER_DIR", "")
    if d:
        return Path(d)
    shared = os.environ.get("NEV_LANE_SHARED", "") or os.environ.get("GENERA_LANE_SHARED", "")
    if shared and Path(shared).is_dir():
        return Path(shared) / "ledger"
    return None


def record(measure_dir):
    d = Path(measure_dir)
    start, end = span(d / "phases.tsv")
    log = d / "run.stamped.log"
    lines = log.read_text(errors="replace").splitlines() if log.exists() else []
    return {
        "format": 1,
        "track": track_name(d),
        "start": start,
        "end": end,
        "cpu_seconds": cpu_seconds(d / "cpu-phases.txt"),
        "models": parse_log(lines),
    }


def main():
    args = sys.argv[1:]
    to_stdout = "--print" in args
    args = [a for a in args if a != "--print"]
    if len(args) != 1:
        print("usage: bench/ledger.py [--print] <measure dir>", file=sys.stderr)
        return 0
    try:
        rec = record(args[0])
    except OSError as e:
        print(f"ledger: not written: cannot read the measurements ({e.strerror})", file=sys.stderr)
        return 0
    line = json.dumps(rec, sort_keys=True, separators=(",", ":"))
    if to_stdout:
        print(line)
        return 0
    if not rec["track"]:
        print("ledger: not written: no track name", file=sys.stderr)
        return 0
    out_dir = ledger_dir()
    if out_dir is None:
        print("ledger: not written: no shared ledger directory (NEV_LANE_SHARED or NEV_LEDGER_DIR)",
              file=sys.stderr)
        return 0
    try:
        out_dir.mkdir(parents=True, exist_ok=True)
        # One write of one line with O_APPEND, so two batches of a track appending at
        # once do not interleave their lines.
        fd = os.open(out_dir / f"{rec['track']}.jsonl", os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0o644)
        try:
            os.write(fd, (line + "\n").encode())
        finally:
            os.close(fd)
    except OSError as e:
        print(f"ledger: not written: {e.strerror}", file=sys.stderr)
        return 0
    print(f"ledger: appended this batch to ledger/{rec['track']}.jsonl", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())