#!/bin/bash
# Reads RustDesk peer ID via GA exec (rustdesk --get-id). Modern RustDesk stores enc_id in TOML, not plain id.
# Install: /usr/local/sbin/blair-sync-rustdesk-via-ga.sh  Cron: every 2 minutes.
# Usage: blair-sync-rustdesk-via-ga.sh [--dry-run]   BLAIR_VERBOSE=1 for detail.
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 binascii
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')
VERBOSE = os.environ.get('BLAIR_VERBOSE', '') in ('1', 'true', 'yes')
VMIDS_ENV = os.environ.get('BLAIR_VMIDS', '').strip()

# RustDesk >=1.2 stores enc_id in TOML, not plain id. Primary: rustdesk.exe --get-id via GA exec.
RUSTDESK_EXE = r'C:\Program Files\RustDesk\rustdesk.exe'
GUEST_REPORT = r'C:\ProgramData\Blair\rustdesk-report.json'


def log(msg: str, *, force: bool = False) -> None:
    if VERBOSE or force:
        print(msg, flush=True)


def update_desc(old: str, new_id: str) -> str:
    old = old.replace('Registry%3A S%3A', 'Registry%3A Z%3A').replace('Registry: S:', 'Registry: Z:')
    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_ga_payload(raw) -> str:
    if raw is None:
        return ''
    if isinstance(raw, bytes):
        return raw.decode('utf-8-sig', errors='replace').strip()
    text = str(raw).strip()
    if not text:
        return ''
    if text[0] in '{[#' or 'rustdesk_id' in text or re.search(r'(?m)^\s*id\s*=', text):
        return text.lstrip('\ufeff')
    if re.fullmatch(r'[\d\r\n\t ]+', text):
        return text.strip()
    blob = ''.join(text.split())
    if not blob:
        return text
    for candidate in (blob, blob + ('=' * ((4 - len(blob) % 4) % 4))):
        try:
            decoded = base64.b64decode(candidate, validate=False)
            return decoded.decode('utf-8-sig', errors='replace').strip()
        except (ValueError, binascii.Error):
            continue
    return text


def clean_exec_text(text: str) -> str:
    if not text:
        return ''
    return re.sub(r'#< CLIXML.*', '', text, flags=re.DOTALL).strip()


def extract_id_from_exec(out: str, err: str) -> str | None:
    combined = '\n'.join(filter(None, (clean_exec_text(out), clean_exec_text(err))))
    for line in combined.splitlines():
        line = line.strip()
        if re.fullmatch(r'\d{6,12}', line):
            return line
    match = re.search(r'\d{6,12}', combined)
    return match.group(0) if match else None


def parse_exec_result(st: dict) -> tuple[str | None, str | None]:
    try:
        exitcode = int(st.get('exitcode', 1))
        out = decode_ga_payload(st.get('out-data'))
        err = decode_ga_payload(st.get('err-data'))
        rid = extract_id_from_exec(out, err)
        if rid:
            return rid, None
        if exitcode != 0:
            brief = clean_exec_text(err) or clean_exec_text(out) or f'guest exit {exitcode}'
            return None, brief[:120] + ('...' if len(brief) > 120 else '')
        return None, f'no id in exec output: {out!r}'
    except Exception as exc:  # noqa: BLE001
        return None, f'exec decode error: {exc}'


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(1)
            continue
        if st.get('exited') != 1:
            time.sleep(1)
            continue
        return parse_exec_result(st)

    return None, 'guest exec timeout'


def parse_id_from_content(text: str, path: str) -> str | None:
    if path.lower().endswith('.json'):
        try:
            rid = str(json.loads(text).get('rustdesk_id', '') or '').strip()
            if re.fullmatch(r'\d{6,12}', rid):
                return rid
        except json.JSONDecodeError:
            pass
        match = re.search(r'"rustdesk_id"\s*:\s*"(\d{6,12})"', text)
        if match:
            return match.group(1)
    for pattern in (
        r'(?m)^\s*id\s*=\s*["\']?(\d{6,12})',
        r'(?m)rustdesk_id\s*=\s*["\']?(\d{6,12})',
    ):
        match = re.search(pattern, text)
        if match:
            return match.group(1)
    return None


def guest_file_read(vmid: int, guest_path: str) -> tuple[str | None, str | None]:
    for path in {guest_path, guest_path.replace('\\', '/')}:
        try:
            proc = subprocess.run(
                [
                    'pvesh', 'get', f'/nodes/{NODE}/qemu/{vmid}/agent/file-read',
                    '--file', path, '--output-format', 'json',
                ],
                capture_output=True, text=True, check=False,
            )
            if proc.returncode != 0:
                continue
            payload = json.loads(proc.stdout)
            data = payload.get('data', payload)
            if not isinstance(data, dict):
                continue
            raw = data.get('content')
            if raw is None:
                continue
            text = decode_ga_payload(raw)
            if text:
                return text, None
        except (json.JSONDecodeError, ValueError, UnicodeError):
            continue
    return None, f'file-read failed for {guest_path}'


def read_rustdesk_id_cli(vmid: int) -> tuple[str | None, str | None]:
    attempts = [
        ([RUSTDESK_EXE, '--get-id'], 'direct'),
        (['cmd.exe', '/c', f'"{RUSTDESK_EXE} --get-id"'], 'cmd'),
        (
            [
                'powershell.exe', '-NoProfile', '-NonInteractive', '-Command',
                "$ProgressPreference='SilentlyContinue'; & 'C:\\Program Files\\RustDesk\\rustdesk.exe' --get-id",
            ],
            'powershell',
        ),
    ]
    errors: list[str] = []
    for command, label in attempts:
        rid, err = guest_exec_text(vmid, command, timeout=45)
        if rid and re.fullmatch(r'\d{6,12}', rid):
            log(f'VM {vmid}: ID {rid} from rustdesk --get-id ({label})')
            return rid, None
        errors.append(f'{label}: {err or "no id"}')
    return None, errors[-1] if errors else 'rustdesk --get-id failed'


def read_rustdesk_id_report(vmid: int) -> tuple[str | None, str | None]:
    text, err = guest_file_read(vmid, GUEST_REPORT)
    if not text:
        return None, err
    rid = parse_id_from_content(text, GUEST_REPORT)
    return (rid, None) if rid else (None, 'no rustdesk_id in report json')


def read_rustdesk_id(vmid: int) -> tuple[str | None, str | None]:
    rid, err = read_rustdesk_id_cli(vmid)
    if rid:
        return rid, None
    log(f'VM {vmid}: cli miss: {err}')

    rid, err = read_rustdesk_id_report(vmid)
    if rid:
        log(f'VM {vmid}: ID {rid} from {GUEST_REPORT}')
        return rid, None

    return None, err or 'no RustDesk ID found'


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}', force=True)
        return 1

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

    registry = load_registry()
    vmids = fleet_vmids(registry)
    if not vmids:
        return 0

    updated: list[int] = []
    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)', force=True)
            continue

        new_id, err = read_rustdesk_id(vmid)
        if not new_id:
            log(f'VM {vmid}: skip ({err})', force=True)
            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}', force=True)

        if DRY_RUN:
            preview = newdesc if len(newdesc) <= 120 else newdesc[:120] + '...'
            log(f'  would set: {preview}', force=True)
            updated.append(vmid)
            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
        save_registry(registry)
        updated.append(vmid)

    if updated and not DRY_RUN:
        log(f'Updated VM(s): {", ".join(str(v) for v in updated)}', force=True)
    elif DRY_RUN and updated:
        log(f'Would update VM(s): {", ".join(str(v) for v in updated)}', force=True)

    return 0


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