#!/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")
    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.
    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"\$ ", timeout=60)
            break
        if i in (1, 2):
            break
        c.sendline("")
    else:
        print("never reached a prompt", file=sys.stderr)
        return 1

    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 ""
    c.sendline("save")
    c.expect(r"# ", timeout=120)
    c.sendline("exit")
    c.expect(r"\$ ", timeout=60)
    c.sendline("exit")
    c.close()

    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())