#!/usr/bin/env python3
"""Prepare/apply/restore fixed PCI hardware for a PVE Q35 VM.

Only adjusts QEMU args. Never stops/starts a VM or restarts its USB tunnel.
Apply/restore require a stopped VM so the next cold start uses the new ACPI.
"""
import argparse
import datetime
import json
from pathlib import Path
import shlex
import subprocess


FIXED_GLOBALS = (
    "ICH9-LPC.acpi-pci-hotplug-with-bridge-support=off",
    "pci-bridge.shpc=off",
    "pcie-root-port.hotplug=off",
)


def fixed_args(original):
    tokens = shlex.split(original)
    additions = []
    for setting in FIXED_GLOBALS:
        prop = setting.split("=", 1)[0]
        for i, token in enumerate(tokens[:-1]):
            if token == "-global" and tokens[i + 1].startswith(prop + "="):
                if tokens[i + 1] != setting:
                    raise ValueError("Conflicting existing setting: " + tokens[i + 1])
                break
        else:
            additions.extend(["-global", setting])
    # Keep the original text, quoting and GPU arguments intact.
    return " ".join(x for x in (original.strip(), shlex.join(additions)) if x)


def command(args):
    return subprocess.run(args, check=True, capture_output=True, text=True,
                          timeout=40).stdout.strip()


def config(vm):
    text = command(["qm", "config", str(vm), "--current", "1"])
    result = {}
    for line in text.splitlines():
        if ": " in line:
            name, value = line.split(": ", 1)
            result[name] = value
    return text, result


def require_stopped(vm):
    if command(["qm", "status", str(vm)]) != "status: stopped":
        raise RuntimeError("VM must be normally shut down first; no live hardware changes made")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["prepare", "apply", "restore"])
    parser.add_argument("--vm", required=True, type=int)
    parser.add_argument("--backup", required=True, type=Path,
                        help="dedicated backup directory, reused for apply/restore")
    args = parser.parse_args()
    if args.vm < 100:
        parser.error("Invalid VM ID")
    text, conf = config(args.vm)
    if "q35" not in conf.get("machine", ""):
        raise RuntimeError("Only an explicitly configured Q35 machine is supported")
    manifest = args.backup / "plan.json"
    if args.action == "prepare":
        if manifest.exists():
            raise RuntimeError("Backup already exists; use a new directory")
        # Pending changes may otherwise be applied unexpectedly at cold start.
        pending = command(["qm", "pending", str(args.vm)])
        if any(not line.startswith("cur ") for line in pending.splitlines() if line):
            raise RuntimeError("VM has pending changes; review them before preparing")
        proposed = fixed_args(conf.get("args", ""))
        args.backup.mkdir(parents=True, exist_ok=True, mode=0o700)
        (args.backup / "vm.conf.before").write_text(text + "\n")
        plan = {"vm": args.vm, "created": datetime.datetime.now().astimezone().isoformat(),
                "original_args": conf.get("args"), "proposed_args": proposed,
                "original_config": conf, "status": "prepared; not applied"}
        manifest.write_text(json.dumps(plan, indent=2) + "\n")
        manifest.chmod(0o600)
        print(json.dumps(plan, indent=2))
        return
    plan = json.loads(manifest.read_text())
    if plan["vm"] != args.vm:
        raise RuntimeError("Backup belongs to another VM")
    require_stopped(args.vm)
    original = plan["original_args"]
    proposed = plan["proposed_args"]
    current = conf.get("args")
    expected = original if args.action == "apply" else proposed
    desired = proposed if args.action == "apply" else original
    if current != expected:
        if current == desired:
            print("Already in requested state")
            return
        raise RuntimeError("QEMU args changed since preparation; refusing to overwrite")
    if args.action == "apply" and conf != plan["original_config"]:
        raise RuntimeError("VM config changed since preparation; prepare a fresh plan")
    if desired is None:
        command(["qm", "set", str(args.vm), "--delete", "args"])
    else:
        command(["qm", "set", str(args.vm), "--args", desired])
    # The script never replaces the full VM config, preserving other settings.
    _, after = config(args.vm)
    if after.get("args") != desired:
        raise RuntimeError("Configuration readback differs; inspect before starting the VM")
    before_other = {k: v for k, v in conf.items() if k != "args"}
    after_other = {k: v for k, v in after.items() if k != "args"}
    if before_other != after_other:
        raise RuntimeError("Other settings changed concurrently; inspect before starting the VM")
    print(json.dumps({"vm": args.vm, "action": args.action,
                      "args": desired, "next_step": "cold start and verify Windows"}))


if __name__ == "__main__":
    main()