#!/usr/bin/env python3
"""Apply VyOS config to a labsim VM over its serial console.

Needed because a freshly installed VyOS comes up holding the same addresses as
its peer, so there is a window where it cannot safely be reached over the
network at all. The console does not care.

  ./console-apply.py --vm labsim-vyos2 --config r2.conf
"""
from __future__ import annotations

import argparse
import sys
import time

import pexpect


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--vm", required=True)
    ap.add_argument("--config", required=True)
    ap.add_argument("--user", default="vyos")
    ap.add_argument("--password", default="vyos")
    # `save` writes config.boot. For WAN work that is dangerous: the resting
    # state must stay `vif 53 disable` on both routers, and saving while a box
    # is master persists the ENABLED state -- so a reboot would have it claim
    # the cloned MAC. Observed in labsim on 2026-09-02.
    ap.add_argument("--no-save", action="store_true",
                    help="commit without saving (leave config.boot untouched)")
    args = ap.parse_args()

    cmds = [l.rstrip() for l in open(args.config)
            if l.strip() and not l.lstrip().startswith("#")]
    print(f"{len(cmds)} commands to apply to {args.vm}", file=sys.stderr)

    c = pexpect.spawn(f"virsh --connect qemu:///system console {args.vm}",
                      timeout=90, encoding="utf-8")
    c.logfile_read = None
    c.sendline("")
    time.sleep(2)
    c.sendline("")

    # Log in. A freshly booted box may still be starting services, so allow a
    # generous window and re-prod the console rather than failing on the first
    # miss.
    #
    # `# ` matters as much as `$ `: a previous run that died mid-config leaves
    # the console sitting in configuration mode, and waiting only for the
    # operational prompt then hangs forever against a perfectly healthy VM.
    in_config = False
    for _ in range(40):
        i = c.expect([r"login:", r"\$ ", r"# ", pexpect.TIMEOUT], timeout=15)
        if i == 0:
            c.sendline(args.user)
            c.expect("Password:", timeout=30)
            c.sendline(args.password)
            c.expect([r"\$ ", r"# "], timeout=60)
            break
        if i == 1:
            break
        if i == 2:
            in_config = True
            break
        c.sendline("")
    else:
        print("never reached a prompt", file=sys.stderr)
        return 1

    if in_config:
        # Drop whatever the previous run left half-built rather than committing
        # a candidate nobody has seen.
        print("console was left in config mode; discarding stale candidate",
              file=sys.stderr)
        c.sendline("discard")
        c.expect(r"# ", timeout=60)
    else:
        c.sendline("configure")
        c.expect(r"# ", timeout=60)

    for cmd in cmds:
        c.sendline(cmd)
        c.expect(r"# ", timeout=60)
        out = c.before or ""
        if "Set failed" in out or "not valid" in out or "Invalid" in out:
            print(f"FAILED: {cmd}\n  {out.strip()[:200]}", file=sys.stderr)

    print("committing...", file=sys.stderr)
    c.sendline("commit")
    c.expect(r"# ", timeout=300)
    commit_out = c.before or ""
    if not args.no_save:
        c.sendline("save")
        c.expect(r"# ", timeout=120)
    # Accept either prompt on the way out. Insisting on `$ ` here hangs against
    # a healthy box -- and worse, leaves the console parked in config mode, so
    # the NEXT run finds a `# ` it was not expecting either. One strict expect
    # turned into two failures.
    c.sendline("exit")
    c.expect([r"\$ ", r"# ", pexpect.TIMEOUT], timeout=60)
    c.sendline("exit")
    c.close(force=True)

    bad = [l for l in commit_out.splitlines()
           if "failed" in l.lower() or "error" in l.lower()]
    if bad:
        print("commit reported:", file=sys.stderr)
        for l in bad[:10]:
            print(f"  {l.strip()}", file=sys.stderr)
        return 1
    print("committed and saved", file=sys.stderr)
    return 0


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