#!/usr/bin/env python3
"""Drive Heretic's interactive Pareto menu non-interactively, and export a chosen trial.

WHY THIS EXISTS
---------------
Heretic finishes a search by printing the Pareto-optimal trials and opening a
`questionary` menu to pick one. `--export-strategy MERGE` chooses HOW to export,
not WHICH trial — the selection is a separate prompt. Run under `nohup`/`</dev/null`
the menu raises `Warning: Input is not a terminal` and the whole run exits 1 with
300 trials of completed search stranded in the study checkpoint.

Nothing is lost when that happens: re-running with the same `--study-checkpoint-dir`
reloads the finished study and goes straight back to the menu. This script gives
that menu a real PTY, reads what it offers, picks by MEASURED VALUE rather than by
position, and answers the follow-up prompts.

WHY NOT --reproduce
-------------------
`reproduce.json` carries only `direction_index` + per-component kernel params. The
published examples are Heretic 1.2.0 and predate `direction_scope`, which 1.4.0
uses to switch between a numeric direction index and per-layer directions. Our
winning trial is scope=0; hand-writing a file whose schema cannot express that
risks silently reproducing a different ablation than the one we measured. Driving
the real menu keeps Heretic's own selection logic authoritative.

SELECTION IS BY VALUE, NOT INDEX
--------------------------------
The menu is ordered by Heretic, not by us, and the Pareto set changes with the
study. Choosing "the Nth entry" would be exactly the positional-inference mistake
that produced four wrong readings during this run. So: parse every menu line for
its refusal count and KL, pick the line whose numbers match the target, and refuse
to guess if no line matches.
"""
from __future__ import annotations

import argparse
import os
import pty
import re
import select
import sys
import time

# Heretic's real row format, confirmed by capturing the live menu:
#     » [Trial 260] Refusals:  8/100, KL divergence: 0.0136
# The count comes AFTER the label. An earlier guess had it before, matched
# nothing, and the driver correctly refused to navigate rather than guess.
ROW = re.compile(r"\[Trial\s+(\d+)\]\s*Refusals:\s*(\d+)\s*/\s*\d+.*?"
                 r"KL\s+divergence:\s*([0-9.]+)", re.I)
ANSI = re.compile(r"\x1b\[[0-9;?]*[a-zA-Z]|\x1b\][^\x07]*\x07|\x1b[=>]")

DOWN, UP, ENTER = b"\x1b[B", b"\x1b[A", b"\r"


def clean(s: str) -> str:
    return ANSI.sub("", s)


class Driver:
    def __init__(self, argv, log_path, timeout=1800, quiet_for=2.5):
        self.argv, self.timeout, self.quiet_for = argv, timeout, quiet_for
        self.log = open(log_path, "w", buffering=1, errors="replace")
        self.buf = ""

    def start(self):
        self.pid, self.fd = pty.fork()
        if self.pid == 0:                      # child
            os.execvp(self.argv[0], self.argv)
        os.set_blocking(self.fd, False)

    def pump(self, until_quiet=None, deadline=None):
        """Read until the child goes quiet for `until_quiet` seconds (or deadline)."""
        until_quiet = self.quiet_for if until_quiet is None else until_quiet
        deadline = deadline or (time.time() + self.timeout)
        last = time.time()
        while time.time() < deadline:
            r, _, _ = select.select([self.fd], [], [], 0.4)
            if r:
                try:
                    chunk = os.read(self.fd, 65536)
                except OSError:
                    break
                if not chunk:
                    break
                text = chunk.decode("utf-8", "replace")
                self.buf += text
                self.log.write(text)
                sys.stdout.write(text)
                sys.stdout.flush()
                last = time.time()
            elif time.time() - last > until_quiet:
                return True
        return False

    def pump_until(self, pattern, deadline_s, label=""):
        """Read until `pattern` appears in the buffer. Returns True if seen.

        Waiting for a MARKER, not for silence. A quiet-based wait cannot survive
        the model-load phase: pulling 52 GB off ZFS and quantizing it pauses for
        longer than any sane quiet threshold, so 'it went quiet' means 'the disk
        stalled', not 'it is ready for input'.
        """
        rx = re.compile(pattern, re.I)
        end = time.time() + deadline_s
        while time.time() < end:
            if rx.search(clean(self.buf)):
                return True
            r, _, _ = select.select([self.fd], [], [], 1.0)
            if r:
                try:
                    chunk = os.read(self.fd, 65536)
                except OSError:
                    break
                if not chunk:
                    break
                text = chunk.decode("utf-8", "replace")
                self.buf += text
                self.log.write(text)
                sys.stdout.write(text)
                sys.stdout.flush()
        print(f"\n[driver] TIMEOUT waiting for {label or pattern!r}", file=sys.stderr)
        return False

    def send(self, data, times=1, pause=0.12):
        for _ in range(times):
            os.write(self.fd, data)
            time.sleep(pause)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--study-checkpoint-dir", required=True)
    ap.add_argument("--out", required=True, help="directory to save the merged model to")
    ap.add_argument("--target-refusals", type=int, required=True)
    ap.add_argument("--target-kl", type=float, required=True)
    ap.add_argument("--kl-tolerance", type=float, default=0.002)
    ap.add_argument("--n-trials", type=int, default=300)
    ap.add_argument("--quantization", default="BNB_4BIT")
    ap.add_argument("--log", default="/tmp/heretic-export-pty.log")
    ap.add_argument("--dry-run", action="store_true",
                    help="show the menu and the chosen row, then quit without saving")
    args = ap.parse_args()

    argv = ["/home/infra-ops/.local/bin/heretic",
            "--model", args.model,
            "--study-checkpoint-dir", args.study_checkpoint_dir,
            "--n-trials", str(args.n_trials),
            "--quantization", args.quantization,
            "--export-strategy", "MERGE"]

    d = Driver(argv, args.log)
    print(f"[driver] spawning: {' '.join(argv)}\n", flush=True)
    d.start()

    # Model load, then a RESUME PROMPT before the Pareto menu:
    #
    #     » Show the results from the previous run
    #       Ignore the previous run and start from scratch   <-- DELETES the checkpoint
    #       Exit program
    #
    # The destructive option is one arrow-key away from the one we want, and it
    # discards every completed trial. So: never send an arrow here. Verify the
    # highlighted default IS the show-results row, then send a bare ENTER. If the
    # default is anything else, abort and let a human look.
    print("\n[driver] waiting for the resume prompt …\n", flush=True)
    d.pump(until_quiet=6.0, deadline=time.time() + 1500)

    if "How would you like to proceed" in clean(d.buf):
        default_row = ""
        for line in clean(d.buf).splitlines():
            if "»" in line:
                default_row = line.strip()
        print(f"[driver] resume prompt detected; highlighted default: {default_row!r}")
        if "Show the results" not in default_row:
            print("[driver] !! the highlighted option is NOT 'Show the results from the "
                  "previous run'. Refusing to press ENTER — the adjacent option deletes "
                  "the study checkpoint. Inspect the PTY log.", file=sys.stderr)
            os.write(d.fd, b"\x03")
            sys.exit(5)
        print("[driver] accepting default with a bare ENTER (no arrows near the destructive option)")
        d.send(ENTER)
        # Model load happens HERE (~1-3 min: 52 GB off ZFS, then 4-bit quantize).
        # Wait for the Pareto banner by name; a quiet-based wait mistakes a disk
        # stall for readiness and bails mid-load.
        print("\n[driver] waiting for the Pareto banner (model load first) …\n", flush=True)
        if not d.pump_until(r"Pareto optimal", 2400, "Pareto banner"):
            os.write(d.fd, b"\x03")
            sys.exit(6)
        # Banner seen; let the menu itself finish rendering.
        d.pump(until_quiet=4.0, deadline=time.time() + 300)

    screen = clean(d.buf)
    # Only the CURRENT menu render matters — the buffer holds every repaint, so
    # parsing the whole thing would stack duplicate rows and wreck the index
    # arithmetic. Take the last block that contains the prompt.
    if "Which trial do you want to use" in screen:
        screen = screen[screen.rindex("Which trial do you want to use"):]
    rows, seen = [], set()
    for line in screen.splitlines():
        m = ROW.search(line)
        if m:
            trial, refus, kl = int(m.group(1)), int(m.group(2)), float(m.group(3))
            if trial in seen:
                continue
            seen.add(trial)
            rows.append((refus, kl, f"[Trial {trial}] {refus}/100 KL {kl:.4f}"))
    print("\n" + "=" * 72)
    print(f"[driver] parsed {len(rows)} candidate menu rows:")
    for r, k, raw in rows:
        print(f"    refusals {r:>3}   KL {k:.4f}    | {raw[:70]}")
    if not rows:
        # The row regex is a guess at Heretic's wording. If it matched nothing,
        # dump the menu verbatim so the format can be read rather than guessed at
        # a second time.
        print("\n[driver] ROW REGEX MATCHED NOTHING — verbatim tail of the screen:")
        tail = [ln for ln in screen.splitlines() if ln.strip()][-45:]
        for ln in tail:
            print(f"    | {ln[:150]}")

    match = [i for i, (r, k, _) in enumerate(rows)
             if r == args.target_refusals and abs(k - args.target_kl) <= args.kl_tolerance]
    if not match:
        print(f"\n[driver] !! no menu row matches {args.target_refusals} refusals @ KL "
              f"{args.target_kl}±{args.kl_tolerance}. NOT guessing a position — "
              f"inspect {args.log} and rerun with corrected targets.", file=sys.stderr)
        os.write(d.fd, b"\x03")
        sys.exit(4)
    idx = match[0]
    print(f"\n[driver] selecting row {idx} -> {rows[idx][2][:70]}")
    if args.dry_run:
        print("[driver] dry-run, sending SIGINT")
        os.write(d.fd, b"\x03")
        sys.exit(0)

    d.send(DOWN, times=idx)          # menu starts highlighted on row 0
    d.send(ENTER)

    # --- post-selection action menu -----------------------------------------
    # This is a questionary SELECT, which ignores typed text (typing only filters
    # in an autocomplete). So it has to be driven by arrows — and therefore by
    # reading the options, not by assuming their order. Heretic re-applies the
    # ablation before showing it ("Resetting model... Abliterating..."), so wait
    # for the prompt itself rather than for a pause.
    print("\n[driver] waiting for the action menu …")
    if not d.pump_until(r"\?\s+What (do you want|would you like)", 1800, "action menu"):
        os.write(d.fd, b"\x03")
        sys.exit(7)
    d.pump(until_quiet=3.0, deadline=time.time() + 120)

    scr = clean(d.buf)
    scr = scr[scr.rindex("?"):] if "?" in scr else scr
    opts = []
    for line in scr.splitlines():
        s = line.strip()
        if s.startswith("»"):
            opts.append((s.lstrip("» ").strip(), True))
        elif s and not s.startswith("?") and len(opts) and len(s) < 90:
            opts.append((s, False))
    print("[driver] action menu options:")
    for i, (t, cur) in enumerate(opts):
        print(f"    {i}{' *' if cur else '  '} {t[:70]}")

    want = [i for i, (t, _) in enumerate(opts) if re.search(r"\bsave\b", t, re.I)]
    if not want:
        print("[driver] !! no option matching 'save' — NOT guessing. Menu dumped above.",
              file=sys.stderr)
        os.write(d.fd, b"\x03")
        sys.exit(8)
    cur = next((i for i, (_, c) in enumerate(opts) if c), 0)
    delta = want[0] - cur
    print(f"[driver] moving {delta:+d} to '{opts[want[0]][0][:50]}' and selecting")
    d.send(DOWN if delta > 0 else UP, times=abs(delta))
    d.send(ENTER)
    d.pump(until_quiet=3.0, deadline=time.time() + 120)

    # Path prompt IS a text input, so typing is correct here.
    print(f"\n[driver] answering path prompt with {args.out}")
    d.send(args.out.encode())
    d.send(ENTER)

    print("\n[driver] saving (this writes ~52 GB, be patient) …\n")
    ok = d.pump(until_quiet=90.0, deadline=time.time() + 5400)
    print(f"\n[driver] {'child quiet' if ok else 'deadline hit'} — see {args.log}")
    os.write(d.fd, b"\x03")


if __name__ == "__main__":
    main()