#!/bin/bash
# Poll RustDesk IDs via QEMU guest agent; update Proxmox VM notes + registry.json.
# Install to /usr/local/sbin on pve. Cron every 5m. No guest reporter or share reports needed.
# Usage: blair-sync-rustdesk-via-ga.sh [--dry-run]
set -euo pipefail

export BLAIR_SHARE="${BLAIR_SHARE:-/tank/blair-share}"
[[ "${1:-}" == "--dry-run" ]] && export BLAIR_DRY_RUN=1

if ! command -v python3 >/dev/null 2>&1; then
  echo 'python3 required on pve' >&2
  exit 1
fi

exec python3 - <<'PY'
import base64
import json
import os
import re
import subprocess
import time
from datetime import date

SHARE = os.environ.get('BLAIR_SHARE', '/tank/blair-share')
REGISTRY = os.path.join(SHARE, 'fleet', 'registry.json')
NODE = os.environ.get('BLAIR_PVE_NODE', 'pve')
DRY_RUN = os.environ.get('BLAIR_DRY_RUN', '') in ('1', 'true', 'yes')
VMIDS_ENV = os.environ.get('BLAIR_VMIDS', '').strip()

RUSTDESK_CFG = (
    r'C:\Windows\ServiceProfiles\LocalService\AppData\Roaming\RustDesk\config\RustDesk2.toml'
)
PS_GET_ID = (
    r"$p='C:\Windows\ServiceProfiles\LocalService\AppData\Roaming\RustDesk\config\RustDesk2.toml'; "
    r"if (-not (Test-Path -LiteralPath $p)) { exit 1 }; "
    r"$t=Get-Content -Raw -LiteralPath $p; "
    r"if ($t -match '(?m)^\s*id\s*=\s*''?(\d{6,12})') { Write-Output $Matches[1]; exit 0 }; "
    r"exit 1"
)


def log(msg: str) -> None:
    print(msg, flush=True)


def update_desc(old: str, new_id: str) -> str:
    if re.search(r'RustDesk\s+\d+', old):
        return re.sub(r'RustDesk\s+\d+', f'RustDesk {new_id}', old)
    if not old:
        return f'RustDesk {new_id}.'
    return f"{old.rstrip('.')}. RustDesk {new_id}."


def load_registry() -> dict:
    with open(REGISTRY, encoding='utf-8') as fh:
        return json.load(fh)


def fleet_vmids(registry: dict) -> list[int]:
    if VMIDS_ENV:
        return [int(x) for x in VMIDS_ENV.split() if x.isdigit()]
    return [int(s['vmid']) for s in registry.get('windows_gpu_seats', []) if s.get('vmid')]


def registry_old_id(registry: dict, vmid: int) -> str:
    for seat in registry.get('windows_gpu_seats', []):
        if seat.get('vmid') == vmid:
            return str(seat.get('rustdesk_id', '') or '')
    return ''


def qm_running(vmid: int) -> bool:
    proc = subprocess.run(['qm', 'status', str(vmid)], capture_output=True, text=True, check=False)
    return 'status: running' in proc.stdout


def guest_ping(vmid: int) -> bool:
    proc = subprocess.run(
        ['qm', 'guest', 'cmd', str(vmid), 'ping'],
        capture_output=True, text=True, check=False,
    )
    return proc.returncode == 0


def decode_exec_payload(st: dict) -> tuple[str, str]:
    out = ''
    err = ''
    if st.get('out-data'):
        out = base64.b64decode(st['out-data']).decode('utf-8', errors='replace').strip()
    if st.get('err-data'):
        err = base64.b64decode(st['err-data']).decode('utf-8', errors='replace').strip()
    return out, err


def parse_exec_result(st: dict) -> tuple[str | None, str | None]:
    exitcode = int(st.get('exitcode', 1))
    out, err = decode_exec_payload(st)
    if exitcode != 0:
        return None, err or f'guest exit {exitcode}'
    match = re.search(r'\d{6,12}', out)
    if not match:
        return None, f'no id in exec output: {out!r}'
    return match.group(0), None


def guest_exec_text(vmid: int, command: list[str], timeout: int = 90) -> tuple[str | None, str | None]:
    proc = subprocess.run(
        ['qm', 'guest', 'exec', str(vmid), '--'] + command,
        capture_output=True, text=True, check=False,
    )
    if not proc.stdout.strip():
        return None, proc.stderr.strip() or f'qm guest exec failed rc={proc.returncode}'
    try:
        data = json.loads(proc.stdout)
    except json.JSONDecodeError:
        return None, f'exec json parse error: {proc.stdout[:160]!r}'

    if data.get('exited') == 1:
        return parse_exec_result(data)

    pid = data.get('pid')
    if not pid:
        return None, f'exec returned no pid: {data!r}'

    deadline = time.time() + timeout
    while time.time() < deadline:
        st_proc = subprocess.run(
            ['qm', 'guest', 'exec-status', str(vmid), str(pid)],
            capture_output=True, text=True, check=False,
        )
        try:
            st = json.loads(st_proc.stdout)
        except json.JSONDecodeError:
            time.sleep(2)
            continue
        if st.get('exited') != 1:
            time.sleep(2)
            continue
        return parse_exec_result(st)

    return None, 'guest exec timeout'


def guest_file_read(vmid: int, guest_path: str) -> tuple[str | None, str | None]:
    proc = subprocess.run(
        ['pvesh', 'get', f'/nodes/{NODE}/qemu/{vmid}/agent/file-read', '--file', guest_path],
        capture_output=True, text=True, check=False,
    )
    if proc.returncode != 0:
        msg = proc.stderr.strip() or proc.stdout.strip() or f'file-read rc={proc.returncode}'
        return None, msg
    try:
        payload = json.loads(proc.stdout)
    except json.JSONDecodeError:
        return None, f'file-read json error: {proc.stdout[:160]!r}'
    data = payload.get('data', payload)
    content_b64 = data.get('content') if isinstance(data, dict) else None
    if not content_b64:
        return None, 'file-read empty content'
    return base64.b64decode(content_b64).decode('utf-8', errors='replace'), None


def parse_rustdesk_id_from_toml(text: str) -> str | None:
    match = re.search(r'(?m)^\s*id\s*=\s*["\']?(\d{6,12})', text)
    return match.group(1) if match else None


def read_rustdesk_id(vmid: int) -> tuple[str | None, str | None]:
    text, err = guest_file_read(vmid, RUSTDESK_CFG)
    if text:
        rid = parse_rustdesk_id_from_toml(text)
        if rid:
            return rid, None
        err = 'id not found in RustDesk2.toml'
    return guest_exec_text(
        vmid,
        ['powershell.exe', '-NoProfile', '-Command', PS_GET_ID],
    )


def qm_description(vmid: int) -> str:
    proc = subprocess.run(['qm', 'config', str(vmid)], capture_output=True, text=True, check=False)
    for line in proc.stdout.splitlines():
        if line.startswith('description:'):
            return line.split(':', 1)[1].strip()
    return ''


def desc_rustdesk_id(desc: str) -> str:
    match = re.search(r'RustDesk\s+(\d{6,12})', desc)
    return match.group(1) if match else ''


def save_registry(registry: dict) -> None:
    tmp = REGISTRY + '.tmp'
    with open(tmp, 'w', encoding='utf-8') as fh:
        json.dump(registry, fh, indent=2)
        fh.write('\n')
    os.replace(tmp, REGISTRY)


def main() -> int:
    if not os.path.isfile(REGISTRY):
        log(f'Missing registry: {REGISTRY}')
        return 1

    if DRY_RUN:
        log('DRY RUN - no qm set or registry writes')

    registry = load_registry()
    vmids = fleet_vmids(registry)
    if not vmids:
        log('No fleet VMIDs found')
        return 0

    registry_dirty = False
    for vmid in vmids:
        if not qm_running(vmid):
            log(f'VM {vmid}: skip (not running)')
            continue
        if not guest_ping(vmid):
            log(f'VM {vmid}: skip (guest agent not responding)')
            continue

        new_id, err = read_rustdesk_id(vmid)
        if not new_id:
            log(f'VM {vmid}: skip (could not read RustDesk ID: {err})')
            continue

        old_registry_id = registry_old_id(registry, vmid)
        current_desc = qm_description(vmid)
        old_desc_id = desc_rustdesk_id(current_desc)

        if new_id == old_registry_id and new_id == old_desc_id:
            log(f'VM {vmid}: OK ({new_id})')
            continue

        newdesc = update_desc(current_desc, new_id)
        log(f'VM {vmid}: {old_registry_id or old_desc_id or "(none)"} -> {new_id}')

        if DRY_RUN:
            preview = newdesc if len(newdesc) <= 120 else newdesc[:120] + '...'
            log(f'  would set description: {preview}')
            continue

        subprocess.run(['qm', 'set', str(vmid), '-description', newdesc], check=True)
        registry['updated'] = date.today().isoformat()
        for seat in registry.get('windows_gpu_seats', []):
            if seat.get('vmid') == vmid:
                seat['rustdesk_id'] = new_id
                registry_dirty = True

    if registry_dirty and not DRY_RUN:
        save_registry(registry)
        log('Updated registry.json')

    return 0


if __name__ == '__main__':
    raise SystemExit(main())
PY