#!/usr/bin/env python3
"""Apply/query/restore power-menu policy via the existing PVE guest agent.

Uses the repository's existing PVE authentication. Never powers off/reboots a VM.
"""
import argparse
import base64
import json
from pathlib import Path
import runpy
import sys
import time


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['apply', 'query', 'restore'])
    parser.add_argument('--vm', type=int, required=True, choices=[100, 101, 102])
    parser.add_argument('--user', help='Required for Linux guests')
    parser.add_argument('--evidence', type=Path, help='Save the guest-agent result as JSON')
    args = parser.parse_args()
    root = Path(__file__).resolve().parent
    api = runpy.run_path(str(root / 'pve_vm.py'))['api']
    ticket, csrf = runpy.run_path(str(root / 'pve_term_exec.py'))['login']()
    base = f'/api2/json/nodes/starhe/qemu/{args.vm}'
    if api(base + '/status/current', ticket=ticket).get('status') != 'running':
        raise SystemExit('Guest is not running; nothing changed')
    if args.vm == 100:
        source = (root / 'set-guest-power-policy.ps1').read_text(encoding='utf-8')
        source = "$ProgressPreference='SilentlyContinue'\n& {\n" + source + "\n} -Action " + args.action.capitalize()
        encoded = base64.b64encode(source.encode('utf-16le')).decode()
        command = ['powershell.exe', '-NoProfile', '-NonInteractive', '-EncodedCommand', encoded]
    else:
        if not args.user:
            parser.error('--user is required for Linux')
        source = (root / 'set-guest-power-policy.py').read_text(encoding='utf-8')
        command = ['python3', '-c', source, args.action, '--user', args.user]
    pid = api(base + '/agent/exec', json_body={'command': command}, ticket=ticket, csrf=csrf)['pid']
    for _ in range(60):
        time.sleep(1)
        result = api(base + f'/agent/exec-status?pid={pid}', ticket=ticket)
        if result.get('exited'):
            report = json.dumps({'vm': args.vm, **result}, ensure_ascii=True)
            if args.evidence:
                args.evidence.parent.mkdir(parents=True, exist_ok=True)
                args.evidence.write_text(report + '\n', encoding='utf-8')
            print(report)
            if result.get('exitcode') or result.get('out-truncated') or result.get('err-truncated'):
                raise SystemExit(1)
            return
    raise SystemExit('Guest operation timed out; query before retrying')


if __name__ == '__main__':
    main()