#!/usr/bin/env python3
"""Drive the labsim VyOS router over its serial console.

Three phases:
  --phase live       wait for the live system and log in
  --phase install    run `install image` unattended
  --phase configure  apply bond0 (LACP) + per-VLAN gateway addresses

The installer prompt list is the same one the bastion's install driver answers
(src/bastion/src/templates/vyos-install.py.ts). Two of them are easy to miss and
both hang forever rather than failing: the reinstall-only "copy data to the new
image?", and "choose two disks for RAID-1 mirroring?" — every RAID prompt
defaults to YES.
"""
from __future__ import annotations

import argparse
import sys
import time

import pexpect

PASSWORD = "vyos"
PROMPT = r"[\$#] $"


def console(vm: str, timeout: int = 60) -> pexpect.spawn:
    c = pexpect.spawn(f"sudo virsh console {vm} --force", encoding="utf-8", timeout=timeout)
    c.expect("Connected to domain", timeout=30)
    return c


def login(c: pexpect.spawn, timeout: int = 300) -> None:
    """Get to a shell prompt, whether we land at a login or an open session."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        c.sendline("")
        i = c.expect(["login:", PROMPT, pexpect.TIMEOUT], timeout=20)
        if i == 0:
            c.sendline("vyos")
            c.expect("assword:", timeout=20)
            c.sendline(PASSWORD)
            j = c.expect([PROMPT, "incorrect", pexpect.TIMEOUT], timeout=30)
            if j == 0:
                return
        elif i == 1:
            return
    raise SystemExit("timed out waiting for a VyOS shell")


def run(c: pexpect.spawn, cmd: str, timeout: int = 60) -> str:
    c.sendline(cmd)
    c.expect(PROMPT, timeout=timeout)
    return c.before or ""


def phase_install(c: pexpect.spawn) -> None:
    """Answer `install image` end to end."""
    rules: list[tuple[str, str]] = [
        (r"Would you like to continue\?", "yes"),
        (r"What would you like to name this image\?", ""),
        (r"Please confirm password for the .vyos. user:", PASSWORD),
        (r"Please enter a password for the .vyos. user:", PASSWORD),
        (r"What console should be used by default", "K"),
        # every RAID variant defaults to YES — decline them all
        (r"Would you like to [^?]*RAID-1 mirroring", "no"),
        (r"Installation will delete all data on (?:the drive|both drives)\. Continue\?", "yes"),
        (r"Which one should be used for installation\?", "/dev/vda"),
        (r"Would you like to use all the free space on the drive\?", "yes"),
        (r"Which file would you like as boot config\?", "1"),
        # reinstall-only; unanswered it blocks on stdin until the world ends
        (r"Would you like to copy data to the new image\?", "yes"),
        (r"From which image would you like to save config information\?", "1"),
    ]
    patterns = [r for r, _ in rules] + [r"The image installed successfully",
                                        r"Unable to install VyOS", pexpect.TIMEOUT]

    c.sendline("install image")
    for _ in range(60):
        i = c.expect(patterns, timeout=180)
        if i < len(rules):
            c.sendline(rules[i][1])
            continue
        if i == len(rules):
            print("  installer: success")
            return
        if i == len(rules) + 1:
            raise SystemExit("installer reported failure")
        raise SystemExit("installer went quiet (unanswered prompt?)")
    raise SystemExit("installer exceeded expected prompt count")


def phase_configure(c: pexpect.spawn, vlans: list[tuple[str, str, str]]) -> None:
    """bond0 over eth0+eth1 with LACP, then a gateway address per VLAN."""
    # Production shape: VLAN 1 (management) is the NATIVE/untagged VLAN on the
    # bond, everything else is a tagged vif. This matters beyond fidelity —
    # LACPDUs are untagged, so a trunk with no native VLAN has nowhere to put
    # them and the bond never negotiates.
    native = [v for v in vlans if v[0] == "1"]
    tagged = [v for v in vlans if v[0] != "1"]

    cmds = [
        "configure",
        "set interfaces bonding bond0 mode '802.3ad'",
        "set interfaces bonding bond0 hash-policy 'layer2+3'",
        "set interfaces bonding bond0 lacp-rate 'fast'",
        "set interfaces bonding bond0 member interface 'eth0'",
        "set interfaces bonding bond0 member interface 'eth1'",
        "set service ssh port '22'",
        "set system login user vyos authentication plaintext-password 'vyos'",
    ]
    for vid, name, prefix in native:
        cmds.append(f"set interfaces bonding bond0 address '{prefix}.1/24'")
        cmds.append(f"set interfaces bonding bond0 description '{name} (native)'")
    for vid, name, prefix in tagged:
        cmds.append(f"set interfaces bonding bond0 vif {vid} address '{prefix}.1/24'")
        cmds.append(f"set interfaces bonding bond0 vif {vid} description '{name}'")
    cmds += ["commit", "save", "exit"]

    for cmd in cmds:
        out = run(c, cmd, timeout=180)
        low = out.lower()
        if "invalid" in low or "syntax error" in low or "commit failed" in low:
            print(f"  !! {cmd}\n{out.strip()[-300:]}")
            raise SystemExit(f"config command rejected: {cmd}")
    print("  config committed and saved")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--vm", required=True)
    ap.add_argument("--phase", required=True, choices=["live", "install", "configure"])
    ap.add_argument("--vlans", default="", help="space separated vid:name:prefix:real entries")
    args = ap.parse_args()

    c = console(args.vm)
    try:
        login(c)
        if args.phase == "live":
            print("  live system reachable")
        elif args.phase == "install":
            phase_install(c)
        else:
            vlans = []
            for entry in args.vlans.split():
                parts = entry.split(":")
                if len(parts) >= 3:
                    vlans.append((parts[0], parts[1], parts[2]))
            if not vlans:
                raise SystemExit("no VLANs passed to configure")
            phase_configure(c, vlans)
        return 0
    finally:
        try:
            c.sendline("")
            c.close(force=True)
        except Exception:
            pass


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