#!/usr/bin/env python3
"""Drive ONE nvflash operation on ONE CMP 170HX through a pty, and answer its confirm prompt only after nvflash
itself has echoed the card we meant (PCI bus) and the versions we expect (current and replacement).

nvflash 5.867 ignores -y and refuses piped stdin, so the answer has to come through a real tty. Feeding a blind
'y' after a timer (the upstream recipe) trusts the adapter index; this reads what nvflash prints first.

  write: flash_one.py --nvflash /path/nvflash --index 7 --bus C2 --rom /path/300w.rom --log /persistent/dir \
                      --cur 92.00.67.00.01 --new 92.00.6D.00.0A
  save:  flash_one.py --nvflash /path/nvflash --index 6 --bus C1 --save /persistent/dir/backup_idx6.rom --log ...

Exit: 0 = success marker seen ("A reboot is required" for a write; file of the right size for a save),
      2 = aborted on purpose (bus or version mismatch: 'n' was sent, nothing written),
      3 = nvflash failed, window already consumed, or timeout.
One Falcon window per card per boot: if this exits 3 with "Falcon"/"Detecting GPU failed", that card needs a
cold cycle before any retry. Never run two of these on the same card in one boot.
"""
import argparse, os, pty, re, select, signal, sys, time

ap = argparse.ArgumentParser()
ap.add_argument("--nvflash", required=True)
ap.add_argument("--index", type=int, required=True)
ap.add_argument("--bus", required=True, help="PCI bus hex as nvidia-smi shows it, e.g. C2")
ap.add_argument("--rom")
ap.add_argument("--save")
ap.add_argument("--log", required=True, help="persistent directory for the typescript")
ap.add_argument("--cur", default="92.00.67.00.01")
ap.add_argument("--new", default="92.00.6D.00.0A")
ap.add_argument("--timeout", type=int, default=900)
ap.add_argument("--quiet", type=int, default=10, help="seconds of silence after both versions echoed = prompt")
ap.add_argument("--trust-index", action="store_true", help="nvflash 5.867 prints the bus only in its Results table AFTER the write: confirm on the two version strings, check the bus afterwards (exit 4 if it differs)")
a = ap.parse_args()
if bool(a.rom) == bool(a.save):
    sys.exit("give exactly one of --rom (write) or --save (backup)")

os.makedirs(a.log, exist_ok=True)
mode = "write" if a.rom else "save"
ts_path = os.path.join(a.log, "nvflash_%s_idx%d_bus%s_%s.log" % (mode, a.index, a.bus, time.strftime("%Y%m%dT%H%M%S")))
ts = open(ts_path, "ab", buffering=0)
def note(msg):
    line = "[flash_one %s] %s\n" % (time.strftime("%H:%M:%S"), msg)
    ts.write(line.encode()); sys.stderr.write(line)

argv = [a.nvflash, "-i", str(a.index)] + (["--save", a.save] if a.save else [a.rom])
note("exec: " + " ".join(argv))
pid, fd = pty.fork()
if pid == 0:
    os.execv(a.nvflash, argv)

bus_re = re.compile(r"\bB:\s*0*%s\b" % re.escape(a.bus.upper()))
out = b""; answered = None; last_rx = time.time(); t0 = time.time(); rc = None
try:
    while True:
        r, _, _ = select.select([fd], [], [], 1.0)
        if r:
            try:
                chunk = os.read(fd, 4096)
            except OSError:
                chunk = b""
            if not chunk:
                break
            out += chunk; ts.write(chunk); last_rx = time.time()
        text = out.decode("utf-8", "replace").upper()
        if mode == "write" and answered is None:
            bus_ok = bool(bus_re.search(text))
            cur_ok = a.cur.upper() in text
            new_ok = a.new.upper() in text
            prompt = bool(re.search(r"PRESS 'Y'|CONFIRM|\(Y/N\)", text))
            settled = cur_ok and new_ok and (time.time() - last_rx) >= a.quiet
            if prompt or settled:
                if (bus_ok or a.trust_index) and cur_ok and new_ok:
                    answered = "y"; note("current %s and replacement %s echoed%s: confirming" % (a.cur, a.new, (", bus %s" % a.bus) if bus_ok else " (bus checked after, --trust-index)"))
                else:
                    answered = "n"; note("ABORT: bus_ok=%s cur_ok=%s new_ok=%s (prompt=%s): sending n" % (bus_ok, cur_ok, new_ok, prompt))
                os.write(fd, (answered + "\r\n").encode())
        if time.time() - t0 > a.timeout:
            note("TIMEOUT after %ds, killing nvflash" % a.timeout); os.kill(pid, signal.SIGKILL); break
finally:
    try:
        _, status = os.waitpid(pid, 0); rc = os.WEXITSTATUS(status) if os.WIFEXITED(status) else 128
    except ChildProcessError:
        rc = -1

text = out.decode("utf-8", "replace")
bad = re.search(r"Falcon In HALT|Detecting GPU failed|system restart might be required|Nothing changed|not accessible", text, re.I)
if mode == "write":
    if answered == "n":
        note("result: aborted by bus/version check (nothing written), nvflash rc=%s, log %s" % (rc, ts_path)); sys.exit(2)
    if re.search(r"A reboot is required", text, re.I) and rc == 0:
        seen = re.findall(r"\bB:\s*([0-9A-Fa-f]{2})\b", text)
        if seen and not any(b.upper() == a.bus.upper() for b in seen):
            note("result: WRITTEN but nvflash reports bus %s, expected %s: WRONG CARD, log %s" % (",".join(seen), a.bus, ts_path)); sys.exit(4)
        note("result: SUCCESS marker seen, rc=0, bus %s, log %s" % (",".join(seen) or "not reported", ts_path)); sys.exit(0)
    note("result: FAILED rc=%s%s, log %s" % (rc, (": " + bad.group(0)) if bad else "", ts_path)); sys.exit(3)
else:
    size = os.path.getsize(a.save) if os.path.exists(a.save) else -1
    if size == 1044480 and open(a.save, "rb").read(4) == b"NVGI":
        note("result: saved %s (%d bytes, NVGI), rc=%s, log %s" % (a.save, size, rc, ts_path)); sys.exit(0)
    note("result: save FAILED size=%d rc=%s%s, log %s" % (size, rc, (": " + bad.group(0)) if bad else "", ts_path)); sys.exit(3)