"""
Togen Vision Windows Agent  (DVI-407 / DVI-438)

Quick-enroll (admin token → active immediately):
    vision_agent.exe --server https://togen.example.com --enroll <admin-token>

Self-enroll (open-enroll secret → pending, admin must approve in the Enroll tab):
    vision_agent.exe --server https://togen.example.com --enroll <open-enroll-secret>

Run (after enrollment; token is stored in Windows Credential Manager):
    vision_agent.exe --server https://togen.example.com

Enroll-only mode (for installers — registers and exits, does not start the check-in loop):
    vision_agent.exe --server https://togen.example.com --enroll <token> --enroll-only
"""

import argparse
import ctypes
import ctypes.wintypes
import json
import logging
import logging.handlers
import os
import platform
import shutil
import socket
import subprocess
import sys
import threading
import time
import uuid
import winreg
from datetime import datetime, timezone
from pathlib import Path

import psutil
import requests

try:
    import keyring
    _HAS_KEYRING = True
except ImportError:
    _HAS_KEYRING = False

try:
    from PIL import Image, ImageGrab
    _HAS_PIL = True
except ImportError:
    _HAS_PIL = False

try:
    import av as _av
    _HAS_AV = True
except ImportError:
    _HAS_AV = False

try:
    import paramiko as _paramiko
    _HAS_PARAMIKO = True
except ImportError:
    _paramiko = None  # type: ignore[assignment]
    _HAS_PARAMIKO = False

_AGENT_VERSION = "3.2.97"
# Suppress console-window flash when spawning scp/ssh/schtasks from a windowless parent.
_NO_WINDOW = getattr(subprocess, "CREATE_NO_WINDOW", 0)
_REG_PATH = r"Software\Togen\Vision"
_REG_MACHINE_ID = "MachineId"
_REG_TOKEN_KEY = "AgentToken"        # fallback when keyring unavailable
_KEYRING_SERVICE = "TogenVisionAgent"
_KEYRING_USERNAME = "agent_token"
_DEFAULT_INTERVAL_S = 60
_LOG_FMT = "%(asctime)s %(levelname)s %(message)s"

# Uninstall / self-destruct (DVI-735)
_INSTALL_DIR = Path(os.environ.get("PROGRAMFILES", r"C:\Program Files")) / "Togen" / "Vision"
_TASK_NAME_CHECKIN = "TogenVisionAgent"
_TASK_NAME_CAPTURE = "TogenVisionAgentCapture"
_LAST_CHECKIN_KEY = "LastCheckinAt"
_SELF_DESTRUCT_DAYS = 30

# --- Capture / rsync constants (Phase 3 / DVI-428) --------------------------
# WAN path (public internet) vs LAN path (192.168.250.x subnet, port 22).
# Override both via env vars TOGEN_RSYNC_HOST / TOGEN_RSYNC_PORT.
_RSYNC_HOST_WAN   = "aiv@165.188.108.55"
_RSYNC_PORT_WAN   = 4423
_RSYNC_HOST_LAN   = "aiv@192.168.250.4"
_RSYNC_PORT_LAN   = 22
_RSYNC_LAN_SUBNET = "192.168.250."
RSYNC_REMOTE_BASE = "footage"          # remote: footage/{machine_id}/
_THUMB_MAX = (320, 200)                # thumbnail bounding box (pixels)
_SEGMENT_DURATION_S = 10              # seconds per rolling H.264 segment (video mode)
_FRAMES_DIR = Path(
    os.environ.get("TOGEN_FRAMES_DIR")
    or Path.home() / ".togen" / "frames"
)
# Machine-level SSH identity (DVI-693): stored in ProgramData so all users
# (local accounts, Azure AD, SYSTEM service) share the same key without per-user
# profile dependency.  Falls back to SSH default key discovery when absent (dev/CI).
_RSYNC_SSH_KEY = Path(os.environ.get("PROGRAMDATA", r"C:\ProgramData")) / "Togen" / "id_togen_vision"

# Machine-wide drop point shared between the SYSTEM check-in process and the
# interactive-session capture process (DVI-683). The check-in process publishes
# the effective capture policy here; the capture process publishes its live
# telemetry (idle/foreground/last-capture) back. ProgramData is readable by both
# SYSTEM and the logged-in user, so neither needs the other's profile hive.
_SHARE_DIR = Path(os.environ.get("PROGRAMDATA", r"C:\ProgramData")) / "Togen" / "Vision"
_POLICY_SHARE = _SHARE_DIR / "policy.json"
_STATE_SHARE = _SHARE_DIR / "capture_state.json"
# Agent log written here; SYSTEM check-in and interactive capture workers share
# the same directory so both streams end up in one place the server can fetch.
_AGENT_LOG_FILE = _SHARE_DIR / "agent.log"
# The interactive --capture-only worker logs to its OWN file: on Windows the
# SYSTEM check-in process holds agent.log open, so a second process's lines were
# silently dropped — which hid scp/push failures from the server log viewer.
_CAPTURE_LOG_FILE = _SHARE_DIR / "capture.log"
_AGENT_LOG_MAX_BYTES = 2 * 1024 * 1024   # 2 MB on-disk; server tails to 200 k
_AGENT_LOG_BACKUP_COUNT = 2
# A capture telemetry record older than this is treated as stale (worker dead /
# no user logged in), so the check-in falls back to local values.
_STATE_STALE_S = 180

# Shared state between the main check-in thread and the capture daemon thread.
_pol_lock = threading.Lock()
_cap_policy = {
    "capture_enabled": False,
    "capture_mode": "screenshot",      # "screenshot" | "video"
    "capture_fps": 1.0,                # video fps; ignored in screenshot mode
    "capture_rate_s": 60,
    "idle_threshold_s": 300,
    "exclusions": [],
    "schedule": {"days": [1, 2, 3, 4, 5], "start": "08:00", "end": "17:00"},
}
_cap_state = {"last_ts": None}  # last successful capture epoch (float)


# ---------------------------------------------------------------------------
# Registry helpers — machine_id
# ---------------------------------------------------------------------------

def _reg_read(hive, path, key):
    try:
        h = winreg.OpenKey(hive, path)
        val, _ = winreg.QueryValueEx(h, key)
        winreg.CloseKey(h)
        return val
    except FileNotFoundError:
        return None


def _reg_write(hive, path, key, value):
    h = winreg.CreateKeyEx(hive, path, access=winreg.KEY_WRITE)
    winreg.SetValueEx(h, key, 0, winreg.REG_SZ, value)
    winreg.CloseKey(h)


def get_machine_id() -> str:
    """Return the stable machine UUID, generating and persisting it on first call.

    Canonical location: HKLM\\Software\\Togen\\Vision\\MachineId.
    Falls back to HKCU when HKLM is not writable (non-admin session).
    """
    # Read: prefer HKLM, fall through to HKCU.
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        val = _reg_read(hive, _REG_PATH, _REG_MACHINE_ID)
        if val:
            return val

    machine_id = str(uuid.uuid4())
    # Write: prefer HKLM, fall back to HKCU on PermissionError.
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        try:
            _reg_write(hive, _REG_PATH, _REG_MACHINE_ID, machine_id)
            return machine_id
        except PermissionError:
            continue

    # No registry write succeeded; machine_id is transient for this run.
    logging.getLogger(__name__).warning(
        "Could not persist machine_id to registry; value is session-only."
    )
    return machine_id


# ---------------------------------------------------------------------------
# Legacy Vision cleanup (DVI-722)
#
# The PowerShell-only AI Vision precursor deployed to C:\Local\ under several
# directory names and registered scheduled tasks outside the TogenVision* prefix.
# This function removes those remnants on first run so machines that were
# enrolled before the Togen-era agent was deployed don't carry stale cruft.
# ---------------------------------------------------------------------------

_LEGACY_VISION_DIRS: list[Path] = [
    Path(r"C:\Local\AI Vision"),
    Path(r"C:\Local\AIVision"),
    Path(r"C:\Local\Togen Vision"),
    Path(r"C:\Local\TogenVision"),
    Path(r"C:\Local\Vision"),
]

# Old scheduled task names used by the PowerShell-era scripts.
# We only remove tasks whose names do NOT start with "TogenVision" (the
# current canonical prefix) to avoid touching our own live tasks.
_LEGACY_TASK_PREFIXES = [
    "AI Vision",
    "AIVision",
    "Vision Agent",
    "VisionAgent",
    "Vision Capture",
    "VisionCapture",
    "Togen Vision Agent",   # old name before the prefix was standardised
]

# Old registry paths outside the current HKLM\Software\Togen\Vision tree.
_LEGACY_REG_KEYS: list[tuple] = [
    (winreg.HKEY_LOCAL_MACHINE, r"Software\AI Vision"),
    (winreg.HKEY_LOCAL_MACHINE, r"Software\AIVision"),
    (winreg.HKEY_CURRENT_USER,  r"Software\AI Vision"),
    (winreg.HKEY_CURRENT_USER,  r"Software\AIVision"),
]


def _cleanup_legacy_vision() -> list[str]:
    """Remove old PowerShell-era AI Vision deployments.

    Checks C:\\Local\\* for known legacy directories, removes old scheduled
    tasks that predate the TogenVision* naming convention, and cleans up
    stale registry keys.  Returns a list of actions taken (for logging).
    Safe to call multiple times — silently skips absent targets.
    """
    log = logging.getLogger(__name__)
    actions: list[str] = []

    # 1. Remove legacy C:\Local directories.
    for d in _LEGACY_VISION_DIRS:
        if d.exists():
            try:
                shutil.rmtree(d)
                msg = f"Removed legacy Vision directory: {d}"
                log.info("Legacy cleanup: %s", msg)
                actions.append(msg)
            except Exception as exc:
                log.warning("Legacy cleanup: could not remove %s: %s", d, exc)

    # 2. Remove legacy scheduled tasks.
    for task_name in _LEGACY_TASK_PREFIXES:
        try:
            result = subprocess.run(
                ["schtasks", "/delete", "/tn", task_name, "/f"],
                capture_output=True, text=True, timeout=15,
                creationflags=_NO_WINDOW,
            )
            if result.returncode == 0:
                msg = f"Removed legacy scheduled task: {task_name}"
                log.info("Legacy cleanup: %s", msg)
                actions.append(msg)
        except Exception:
            pass

    # 3. Remove legacy registry keys.
    for hive, path in _LEGACY_REG_KEYS:
        try:
            winreg.DeleteKey(hive, path)
            msg = f"Removed legacy registry key: {path}"
            log.info("Legacy cleanup: %s", msg)
            actions.append(msg)
        except FileNotFoundError:
            pass
        except Exception as exc:
            log.warning("Legacy cleanup: could not remove registry key %s: %s", path, exc)

    if actions:
        log.info("Legacy Vision cleanup complete: %d item(s) removed.", len(actions))
    else:
        log.debug("Legacy Vision cleanup: nothing to remove.")

    return actions


# ---------------------------------------------------------------------------
# Self-uninstall / self-destruct (DVI-735)
# ---------------------------------------------------------------------------

def _get_last_checkin_at() -> float | None:
    """Read epoch of last successful check-in from registry."""
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        val = _reg_read(hive, _REG_PATH, _LAST_CHECKIN_KEY)
        if val:
            try:
                return float(val)
            except (ValueError, TypeError):
                pass
    return None


def _set_last_checkin_at(ts: float):
    """Persist epoch of latest successful check-in to registry."""
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        try:
            _reg_write(hive, _REG_PATH, _LAST_CHECKIN_KEY, str(ts))
            return
        except PermissionError:
            continue


def _self_uninstall(reason: str):
    """Uninstall the Vision Agent cleanly.

    1. Stops and deletes the TogenVisionAgent and TogenVisionAgentCapture
       scheduled tasks (prevents restarts).
    2. Removes the agent token from Windows Credential Manager and HKLM/HKCU
       registry (HKLM\\Software\\Togen\\Vision).
    3. Launches a detached PowerShell script (no console window) that waits
       3 seconds then removes the install dir and ProgramData dir.  The delay
       allows this process to exit before its own files are removed.
    4. Calls sys.exit(0).
    """
    log = logging.getLogger(__name__)
    log.warning("Self-uninstall triggered: %s", reason)

    # 1. Stop then delete scheduled tasks (capture first, then check-in).
    for task_name in (_TASK_NAME_CAPTURE, _TASK_NAME_CHECKIN):
        for sub_cmd in ("/end", "/delete"):
            try:
                subprocess.run(
                    ["schtasks", sub_cmd, "/tn", task_name, "/f"],
                    capture_output=True, timeout=15,
                    creationflags=_NO_WINDOW,
                )
            except Exception:
                pass

    # 2. Remove stored agent token from Windows Credential Manager.
    if _HAS_KEYRING:
        try:
            keyring.delete_password(_KEYRING_SERVICE, _KEYRING_USERNAME)
        except Exception:
            pass

    # 3. Remove HKLM and HKCU registry key (machine_id, token, last-checkin).
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        try:
            winreg.DeleteKey(hive, _REG_PATH)
        except FileNotFoundError:
            pass
        except Exception as exc:
            log.debug("Could not remove registry key: %s", exc)

    # 4. Launch a detached PowerShell cleanup script to remove the install dir
    #    and shared data dir after this process exits (can't delete self while running).
    install_dir = str(_INSTALL_DIR).replace("'", "''")
    data_dir = str(_SHARE_DIR).replace("'", "''")
    ps_script = (
        "Start-Sleep -Seconds 3\n"
        f"Remove-Item -Path '{install_dir}' -Recurse -Force -ErrorAction SilentlyContinue\n"
        f"Remove-Item -Path '{data_dir}' -Recurse -Force -ErrorAction SilentlyContinue\n"
    )
    try:
        import tempfile
        tmp = tempfile.NamedTemporaryFile(
            mode="w", suffix=".ps1", prefix="togen-uninstall-", delete=False, encoding="utf-8",
        )
        tmp.write(ps_script)
        tmp.close()
        ps_exe = (
            shutil.which("powershell.exe")
            or r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe"
        )
        subprocess.Popen(
            [ps_exe, "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden",
             "-NonInteractive", "-File", tmp.name],
            creationflags=subprocess.CREATE_NO_WINDOW | subprocess.CREATE_NEW_PROCESS_GROUP,
        )
        log.info("Uninstall cleanup script launched: %s", tmp.name)
    except Exception as exc:
        log.warning("Could not launch cleanup script: %s", exc)

    log.info("Self-uninstall complete — exiting.")
    sys.exit(0)


def _check_self_destruct_timeout():
    """Self-destruct if the agent has not reached Togen for more than 30 days.

    Writes the initial timestamp on first call so the clock starts cleanly
    on a fresh install rather than triggering immediately.
    """
    log = logging.getLogger(__name__)
    last = _get_last_checkin_at()
    if last is None:
        _set_last_checkin_at(time.time())
        log.debug("Self-destruct timer initialised (no prior check-in recorded).")
        return
    age_days = (time.time() - last) / 86400.0
    if age_days > _SELF_DESTRUCT_DAYS:
        log.warning(
            "Last successful check-in was %.1f days ago (threshold: %d days) — "
            "initiating self-uninstall.",
            age_days, _SELF_DESTRUCT_DAYS,
        )
        _self_uninstall(
            f"no successful check-in for {age_days:.1f} days (>{_SELF_DESTRUCT_DAYS}-day threshold)"
        )


# ---------------------------------------------------------------------------
# Credential storage — agent_token (Windows Credential Manager via keyring)
# ---------------------------------------------------------------------------

def _store_agent_token(token: str):
    if _HAS_KEYRING:
        try:
            keyring.set_password(_KEYRING_SERVICE, _KEYRING_USERNAME, token)
        except Exception:
            pass
    # Write to HKLM so SYSTEM (scheduled task at startup) can read it back.
    # Fall back to HKCU if HKLM write fails (non-admin context).
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        try:
            _reg_write(hive, _REG_PATH, _REG_TOKEN_KEY, token)
            return
        except PermissionError:
            continue


def _load_agent_token() -> str | None:
    if _HAS_KEYRING:
        try:
            val = keyring.get_password(_KEYRING_SERVICE, _KEYRING_USERNAME)
            if val:
                return val
        except Exception:
            pass
    # Check HKLM first (readable by SYSTEM), then fall back to HKCU.
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        val = _reg_read(hive, _REG_PATH, _REG_TOKEN_KEY)
        if val:
            return val
    return None


# ---------------------------------------------------------------------------
# Capture policy helpers (Phase 3 / DVI-428)
# ---------------------------------------------------------------------------

def _apply_config_delta(delta: dict):
    """Overwrite shared policy fields from a check-in response config_delta."""
    with _pol_lock:
        for k in ("capture_enabled", "capture_mode", "capture_fps",
                  "capture_rate_s", "idle_threshold_s", "exclusions", "schedule"):
            if k in delta:
                _cap_policy[k] = delta[k]


def _js_weekday(dt: datetime) -> int:
    """Convert Python weekday (Mon=0..Sun=6) to JS/policy weekday (Sun=0..Sat=6)."""
    return (dt.weekday() + 1) % 7


def _in_schedule(sched: dict) -> bool:
    """Return True if the current local time falls within the policy schedule."""
    now = datetime.now()
    allowed_days = sched.get("days")
    if allowed_days is not None and _js_weekday(now) not in allowed_days:
        return False
    start = sched.get("start", "00:00")
    end = sched.get("end", "23:59")
    now_hhmm = now.strftime("%H:%M")
    return start <= now_hhmm <= end


def _is_excluded(exclusions: list, foreground_app: str | None) -> bool:
    """Return True if the foreground app title matches any exclusion entry."""
    if not exclusions or not foreground_app:
        return False
    app_lower = foreground_app.lower()
    return any(e.lower() in app_lower for e in exclusions if e)


def _should_capture(policy: dict, idle_seconds: int,
                    foreground_app: str | None) -> bool:
    """Evaluate all policy gates; return True only when capture is appropriate."""
    if not policy.get("capture_enabled"):
        return False
    if idle_seconds >= policy.get("idle_threshold_s", 300):
        return False
    sched = policy.get("schedule") or {}
    if sched and not _in_schedule(sched):
        return False
    if _is_excluded(policy.get("exclusions") or [], foreground_app):
        return False
    return True


# ---------------------------------------------------------------------------
# Session detection + cross-process policy/state sharing (DVI-683)
# ---------------------------------------------------------------------------

def _current_session_id():
    """Windows session id of this process (0 == non-interactive service/Session 0)."""
    try:
        pid = ctypes.windll.kernel32.GetCurrentProcessId()
        sid = ctypes.wintypes.DWORD()
        if ctypes.windll.kernel32.ProcessIdToSessionId(pid, ctypes.byref(sid)):
            return int(sid.value)
    except Exception:
        pass
    return None


def _is_session_zero():
    """True when running in Session 0 (the SYSTEM service session) — cannot see
    the interactive desktop, so screen capture must run elsewhere."""
    return _current_session_id() == 0


def _active_console_session():
    """Session id currently attached to the physical console, or None if no one is
    logged in (0xFFFFFFFF). Used by the SYSTEM process to know whether an
    interactive capture worker can run."""
    try:
        sid = ctypes.windll.kernel32.WTSGetActiveConsoleSessionId()
        return None if sid in (0xFFFFFFFF, -1) else int(sid)
    except Exception:
        return None


def _write_shared_policy(policy: dict):
    """SYSTEM check-in process publishes the effective capture policy for the
    interactive capture worker to consume."""
    try:
        _SHARE_DIR.mkdir(parents=True, exist_ok=True)
        _POLICY_SHARE.write_text(json.dumps(policy))
    except OSError as exc:
        logging.getLogger(__name__).debug("Could not write shared policy: %s", exc)


def _read_shared_policy() -> dict | None:
    """Capture worker reads the policy published by the check-in process."""
    try:
        return json.loads(_POLICY_SHARE.read_text())
    except (OSError, json.JSONDecodeError):
        return None


def _patch_shared_policy(updates: dict):
    """Merge `updates` into the shared policy file without overwriting other keys."""
    try:
        existing = _read_shared_policy() or {}
        existing.update(updates)
        _write_shared_policy(existing)
    except Exception as exc:
        logging.getLogger(__name__).debug("Could not patch shared policy: %s", exc)


def _write_shared_state(idle_seconds, foreground_app, last_capture_ts):
    """Capture worker publishes live, session-correct telemetry for the check-in
    process to forward to the server (the SYSTEM process sees only Session-0
    values, which are meaningless for an interactive desktop)."""
    try:
        _SHARE_DIR.mkdir(parents=True, exist_ok=True)
        _STATE_SHARE.write_text(json.dumps({
            "idle_seconds": idle_seconds,
            "foreground_app": foreground_app,
            "last_capture_ts": last_capture_ts,
            "updated_at": time.time(),
        }))
    except OSError as exc:
        logging.getLogger(__name__).debug("Could not write shared state: %s", exc)


def _read_shared_state() -> dict | None:
    """Check-in process reads fresh worker telemetry; None if absent or stale."""
    try:
        st = json.loads(_STATE_SHARE.read_text())
    except (OSError, json.JSONDecodeError):
        return None
    if time.time() - float(st.get("updated_at", 0)) > _STATE_STALE_S:
        return None
    return st


# ---------------------------------------------------------------------------
# Screen capture (Phase 3 / DVI-428)
# ---------------------------------------------------------------------------

def _frame_files(machine_dir: Path, ts: float) -> list[Path]:
    """The full + thumbnail files written by _capture_frame for timestamp ts."""
    t = int(ts)
    return [machine_dir / f"{t}.jpg", machine_dir / f"{t}.thumb.jpg"]


def _capture_frame(machine_id: str, frames_dir: Path) -> float | None:
    """Screenshot + thumbnail. Returns epoch timestamp on success, None on failure.

    File layout (matches togen/app.py _vision_list_frames expectations):
      frames_dir/{machine_id}/{ts}.jpg          — full image
      frames_dir/{machine_id}/{ts}.thumb.jpg    — low-res companion
    """
    if not _HAS_PIL:
        logging.getLogger(__name__).warning(
            "Pillow not installed — capture skipped. Install: pip install Pillow"
        )
        return None
    log = logging.getLogger(__name__)
    ts = int(time.time())
    machine_dir = frames_dir / machine_id
    machine_dir.mkdir(parents=True, exist_ok=True)
    try:
        img = ImageGrab.grab(all_screens=True)
        full_path = machine_dir / f"{ts}.jpg"
        img.save(str(full_path), "JPEG", quality=85, optimize=True)
        thumb = img.copy()
        thumb.thumbnail(_THUMB_MAX, Image.LANCZOS)
        thumb_path = machine_dir / f"{ts}.thumb.jpg"
        thumb.save(str(thumb_path), "JPEG", quality=75, optimize=True)
        log.debug("Frame captured: %s (%dx%d → %dx%d thumb)",
                  full_path.name, img.width, img.height,
                  thumb.width, thumb.height)
        return float(ts)
    except Exception as exc:
        log.warning("Screen capture failed: %s", exc)
        return None


# ---------------------------------------------------------------------------
# Video capture — H.264 segment encoding (DVI-730 / DVI-727 Phase 1)
# ---------------------------------------------------------------------------

def _encode_video_segment(
    frames: list,
    ts: int,
    fps: float,
    machine_dir: "Path",
) -> "tuple[Path | None, Path | None, Path | None]":
    """Encode a list of PIL Images → H.264 .mp4 + poster .jpg + thumb.

    Writes to a .part temp file, then atomically renames to avoid the
    server-side pull fetching a partial segment.

    Returns (seg_path, poster_path, thumb_path) on success;
    (None, None, None) on any failure.
    """
    if not (_HAS_AV and _HAS_PIL and frames):
        return None, None, None
    import fractions
    log = logging.getLogger(__name__)
    tmp_seg = machine_dir / f"{ts}_vid.part"
    seg_path = machine_dir / f"{ts}.mp4"
    poster_path = machine_dir / f"{ts}.jpg"
    thumb_path = machine_dir / f"{ts}.thumb.jpg"
    try:
        w, h = frames[0].size
        # PyAV requires a Fraction (not float) for the stream rate.
        rate = fractions.Fraction(fps).limit_denominator(1000)
        container = _av.open(str(tmp_seg), mode="w", format="mp4")
        stream = container.add_stream("libx264", rate=rate)
        stream.width = w
        stream.height = h
        stream.pix_fmt = "yuv420p"
        stream.options = {"crf": "23", "preset": "veryfast", "tune": "zerolatency"}
        for i, img in enumerate(frames):
            vf = _av.VideoFrame.from_image(img)
            vf.pts = i
            for pkt in stream.encode(vf):
                container.mux(pkt)
        for pkt in stream.encode():
            container.mux(pkt)
        container.close()
        tmp_seg.rename(seg_path)
        # Poster frame (first capture) + thumbnail
        poster_img = frames[0]
        poster_img.save(str(poster_path), "JPEG", quality=85, optimize=True)
        thumb = poster_img.copy()
        thumb.thumbnail(_THUMB_MAX, Image.LANCZOS)
        thumb.save(str(thumb_path), "JPEG", quality=75, optimize=True)
        log.debug("Segment encoded: %s (%d frames, %d B)",
                  seg_path.name, len(frames), seg_path.stat().st_size)
        return seg_path, poster_path, thumb_path
    except Exception as exc:
        log.warning("Video encode failed: %s", exc)
        try:
            tmp_seg.unlink(missing_ok=True)
        except Exception:
            pass
        return None, None, None


def _capture_video_segment(
    machine_id: str,
    frames_dir: "Path",
    pol: dict,
    fps: float,
) -> "float | None":
    """Collect _SEGMENT_DURATION_S seconds of frames at fps, encode, and send.

    Gating (idle / exclusion / schedule) is re-evaluated each frame so capture
    stops promptly if the user leaves or an excluded app gains focus mid-segment.

    Returns the segment start epoch (float) on success; None on failure or if
    no frames were collected.
    """
    log = logging.getLogger(__name__)
    seg_start = int(time.time())
    frame_interval_s = 1.0 / max(fps, 0.01)
    deadline = time.monotonic() + _SEGMENT_DURATION_S
    frames = []
    machine_dir = frames_dir / machine_id
    machine_dir.mkdir(parents=True, exist_ok=True)

    while time.monotonic() < deadline:
        t0 = time.monotonic()
        idle = _get_idle_seconds()
        fapp = _get_foreground_app()
        if not _should_capture(pol, idle, fapp):
            log.debug("Video capture gated mid-segment (idle=%ds); stopping early.", idle)
            break
        try:
            img = ImageGrab.grab(all_screens=True)
            frames.append(img)
        except Exception as exc:
            log.debug("Screen grab failed: %s", exc)
        elapsed = time.monotonic() - t0
        wait = frame_interval_s - elapsed
        if wait > 0:
            time.sleep(wait)

    if not frames:
        return None

    seg_path, poster_path, thumb_path = _encode_video_segment(frames, seg_start, fps, machine_dir)
    if seg_path is None:
        return None

    files_to_send = [f for f in (seg_path, poster_path, thumb_path) if f is not None]
    # Use the same paramiko→scp→rsync priority order as screenshot mode (DVI-804).
    # The previous scp-only path always failed because _ensure_remote_dir runs a shell
    # command on the remote aiv account, which is rsync/SFTP-restricted with no shell.
    if files_to_send:
        _send_frames(machine_dir, machine_id, files_to_send)

    achieved_fps = len(frames) / _SEGMENT_DURATION_S
    log.info("Video segment sent: %s  frames=%d  achieved_fps=%.2f  size=%d B",
             seg_path.name, len(frames), achieved_fps, seg_path.stat().st_size)
    return float(seg_start)


# ---------------------------------------------------------------------------
# Rsync send (Phase 3 / DVI-428)
# ---------------------------------------------------------------------------

def _rsync_endpoint() -> tuple[str, int]:
    """Return (user@host, port) for rsync.

    Prefers the LAN path (192.168.250.4:22) when this machine has a 192.168.250.x
    address; otherwise falls back to the WAN path (165.188.108.55:4423).
    Either leg can be overridden via TOGEN_RSYNC_HOST / TOGEN_RSYNC_PORT env vars.
    """
    host_env = os.environ.get("TOGEN_RSYNC_HOST")
    port_env = os.environ.get("TOGEN_RSYNC_PORT")
    if host_env:
        return host_env, int(port_env or _RSYNC_PORT_WAN)
    try:
        for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET):
            if info[4][0].startswith(_RSYNC_LAN_SUBNET):
                return _RSYNC_HOST_LAN, _RSYNC_PORT_LAN
    except Exception:
        pass
    return _RSYNC_HOST_WAN, _RSYNC_PORT_WAN


def _rsync_frames(local_machine_dir: Path, machine_id: str) -> bool:
    """Rsync captured frames to the remote footage store via SSH.

    Remote path: {RSYNC_REMOTE_BASE}/{machine_id}/ on the rsync server.
    Requires rsync on PATH and an SSH key authorised for aiv@<rsync-host>.
    Falls back gracefully when rsync is missing (dev/testing environments).
    """
    log = logging.getLogger(__name__)
    if not local_machine_dir.is_dir():
        return False
    rsync_host, rsync_port = _rsync_endpoint()
    # Trailing slash = send contents, not the directory itself.
    src = str(local_machine_dir).rstrip("/\\") + "/"
    dst = f"{rsync_host}:{RSYNC_REMOTE_BASE}/{machine_id}/"
    _id_flag = f" -i {_RSYNC_SSH_KEY}" if _RSYNC_SSH_KEY.exists() else ""
    cmd = [
        "rsync", "-az", "--timeout=30",
        "-e", f"ssh -p {rsync_port}{_id_flag} -o StrictHostKeyChecking=no -o BatchMode=yes",
        src, dst,
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=90,
                                creationflags=_NO_WINDOW)
        if result.returncode == 0:
            log.debug("Rsync OK → %s", dst)
            return True
        log.warning("Rsync failed (rc=%d): %s",
                    result.returncode, (result.stderr or result.stdout).strip())
        return False
    except subprocess.TimeoutExpired:
        log.warning("Rsync timed out after 90s")
        return False
    except FileNotFoundError:
        log.warning("rsync not found on PATH — footage send skipped")
        return False


# ---------------------------------------------------------------------------
# paramiko SFTP send (DVI-799 — no system OpenSSH required)
#
# On most Windows 10/11 installs the OpenSSH Client optional feature is NOT
# installed by default, so scp.exe/ssh.exe are absent and the scp path below
# silently falls through to rsync, which is also missing.  Paramiko is a
# pure-Python SSH2 implementation bundled inside the .exe by PyInstaller, so
# it works on every box regardless of optional-feature state.
# ---------------------------------------------------------------------------

def _sftp_frames(files: list[Path], machine_id: str) -> bool:
    """Send frame files via paramiko SFTP (bundled SSH2, no system OpenSSH needed)."""
    log = logging.getLogger(__name__)
    if not _HAS_PARAMIKO:
        return False
    files = [f for f in files if f.is_file()]
    if not files or not _RSYNC_SSH_KEY.exists():
        return False

    raw_host, port = _rsync_endpoint()
    if "@" in raw_host:
        username, hostname = raw_host.split("@", 1)
    else:
        username, hostname = "aiv", raw_host

    try:
        pkey = _paramiko.Ed25519Key.from_private_key_file(str(_RSYNC_SSH_KEY))
    except Exception as exc:
        log.warning("sftp: key load failed (%s): %s", _RSYNC_SSH_KEY, exc)
        return False

    transport = None
    try:
        transport = _paramiko.Transport((hostname, port))
        transport.connect(username=username, pkey=pkey)
        sftp = _paramiko.SFTPClient.from_transport(transport)

        remote_dir = f"{RSYNC_REMOTE_BASE}/{machine_id}"
        # Ensure parent and machine dir exist (server may not have them yet).
        for d in (RSYNC_REMOTE_BASE, remote_dir):
            try:
                sftp.stat(d)
            except IOError:
                try:
                    sftp.mkdir(d)
                except IOError:
                    pass  # race — already created by another worker

        sent = 0
        for f in files:
            try:
                sftp.put(str(f), f"{remote_dir}/{f.name}")
                sent += 1
            except Exception as put_exc:
                log.warning("sftp: put %s failed: %s", f.name, put_exc)

        sftp.close()
        if sent > 0:
            log.debug("sftp OK → %s (%d/%d files)", remote_dir, sent, len(files))
            return True
        log.warning("sftp: 0/%d files sent to %s", len(files), remote_dir)
        return False

    except Exception as exc:
        log.warning("sftp failed: %s", exc)
        return False
    finally:
        if transport is not None:
            try:
                transport.close()
            except Exception:
                pass


# ---------------------------------------------------------------------------
# scp send (DVI-683 — WDAC/Smart App Control safe path)
#
# DESKTOP-UPG92U7 enforces Smart App Control (WDAC base policy
# {0283ac0f-fff1-49ae-ada1-8a933130cad6}, ISG). It blocks the unsigned
# chocolatey rsync.exe at image load (CodeIntegrity Event 3077), so rsync can
# never run there. OpenSSH scp.exe/ssh.exe ship Microsoft-signed and are always
# trusted, so we transfer via scp by preference and fall back to rsync only when
# scp is unavailable (e.g. non-Windows dev/CI hosts).
# ---------------------------------------------------------------------------

_SCP_EXE = shutil.which("scp")
_SSH_EXE = shutil.which("ssh")
_remote_dirs_ready: set[str] = set()


def _ssh_id_args() -> list[str]:
    return ["-i", str(_RSYNC_SSH_KEY)] if _RSYNC_SSH_KEY.exists() else []


def _ensure_remote_dir(machine_id: str, host: str, port: int) -> bool:
    """scp does not create remote directories; ensure footage/<machine_id> exists."""
    if machine_id in _remote_dirs_ready or not _SSH_EXE:
        return machine_id in _remote_dirs_ready
    cmd = [
        _SSH_EXE, "-p", str(port), *_ssh_id_args(),
        "-o", "StrictHostKeyChecking=no", "-o", "BatchMode=yes",
        "-o", "ConnectTimeout=15", host,
        f"mkdir -p {RSYNC_REMOTE_BASE}/{machine_id}",
    ]
    try:
        if subprocess.run(cmd, capture_output=True, text=True, timeout=30,
                          creationflags=_NO_WINDOW).returncode == 0:
            _remote_dirs_ready.add(machine_id)
            return True
    except (subprocess.TimeoutExpired, OSError):
        pass
    return False


def _scp_frames(files: list[Path], machine_id: str) -> bool:
    """Send specific frame files to footage/<machine_id>/ via Microsoft-signed scp."""
    log = logging.getLogger(__name__)
    files = [f for f in files if f.is_file()]
    if not files or not (_SCP_EXE and _SSH_EXE):
        return False
    host, port = _rsync_endpoint()
    if not _ensure_remote_dir(machine_id, host, port):
        log.warning("scp: could not ensure remote dir footage/%s", machine_id)
        return False
    dst = f"{host}:{RSYNC_REMOTE_BASE}/{machine_id}/"
    # -O forces the legacy SCP/rcp transfer protocol (needs only a remote shell)
    # instead of the modern SFTP protocol; the footage server runs an rsync-over-ssh
    # shell account that may not expose the sftp subsystem, where SFTP-mode scp hangs.
    cmd = [
        _SCP_EXE, "-O", "-p", "-P", str(port), *_ssh_id_args(),
        "-o", "StrictHostKeyChecking=no", "-o", "BatchMode=yes",
        "-o", "ConnectTimeout=15",
        *[str(f) for f in files], dst,
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=90,
                                creationflags=_NO_WINDOW)
        if result.returncode == 0:
            log.debug("scp OK → %s (%d files)", dst, len(files))
            return True
        log.warning("scp failed (rc=%d): %s",
                    result.returncode, (result.stderr or result.stdout).strip())
        return False
    except subprocess.TimeoutExpired:
        log.warning("scp timed out after 90s")
        return False
    except OSError as exc:
        log.warning("scp invocation error: %s", exc)
        return False


def _send_frames(machine_dir: Path, machine_id: str, new_files: list[Path] | None) -> bool:
    """Transfer captured frames.

    Priority: paramiko SFTP (bundled, no system OpenSSH) → scp (WDAC-signed) → rsync.
    new_files: files from the latest capture (full + thumb); paramiko/scp send only
    those; rsync sends the whole directory incrementally.
    """
    if _HAS_PARAMIKO and new_files:
        if _sftp_frames(new_files, machine_id):
            return True
    if _SCP_EXE and _SSH_EXE and new_files:
        if _scp_frames(new_files, machine_id):
            return True
    return _rsync_frames(machine_dir, machine_id)


# ---------------------------------------------------------------------------
# Capture daemon thread (Phase 3 / DVI-428)
# ---------------------------------------------------------------------------

def _run_capture_thread(machine_id: str, frames_dir: Path):
    """Daemon thread: capture frames and encode per the active policy.

    Used when the agent runs in an interactive session directly (not via the
    two-task setup). Supports both screenshot and video (DVI-730) modes.
    """
    log = logging.getLogger(__name__)
    log.info("Capture thread started (machine_id=%s)", machine_id)
    while True:
        with _pol_lock:
            pol = dict(_cap_policy)
        mode = pol.get("capture_mode", "screenshot")
        cap_enabled = pol.get("capture_enabled", False)
        idle = _get_idle_seconds()
        fapp = _get_foreground_app()
        if cap_enabled and _should_capture(pol, idle, fapp):
            if mode == "video" and _HAS_AV:
                fps = max(0.1, min(30.0, float(pol.get("capture_fps", 1.0))))
                new_ts = _capture_video_segment(machine_id, frames_dir, pol, fps)
                if new_ts is not None:
                    with _pol_lock:
                        _cap_state["last_ts"] = new_ts
                # segment collection provides its own pacing; loop immediately
                continue
            else:
                ts = _capture_frame(machine_id, frames_dir)
                if ts is not None:
                    with _pol_lock:
                        _cap_state["last_ts"] = ts
                    machine_dir = frames_dir / machine_id
                    _send_frames(machine_dir, machine_id, _frame_files(machine_dir, ts))
        rate = max(10, pol.get("capture_rate_s", 60))
        time.sleep(rate)


def _worker_upload_capture_log(base_url: str | None, token: str | None) -> None:
    """Push the worker's own capture.log to the server log store (DVI-798).

    The SYSTEM check-in process owns the upload_log command, but on Windows it
    holds agent.log open and never sees this worker's capture.log — so push
    failures were invisible server-side. The worker self-uploads its log here so
    the admin log viewer can read scp/ssh stderr without an exe rebuild.
    """
    if not base_url or not token or not _CAPTURE_LOG_FILE.exists():
        return
    try:
        text = _CAPTURE_LOG_FILE.read_text(encoding="utf-8", errors="replace")
        if len(text) > 200_000:
            text = text[-200_000:]
        requests.post(
            f"{base_url.rstrip('/')}/vision/agent-log",
            json={"file": _CAPTURE_LOG_FILE.name, "content": text},
            headers={"Authorization": f"Bearer {token}"},
            timeout=30,
        )
    except Exception:
        pass


def run_capture_only(machine_id: str, frames_dir: Path, base_url: str | None = None):
    """Capture worker entry point (DVI-683 / DVI-730).

    Runs in the interactive console session (launched by an AtLogOn scheduled
    task) so screen capture, idle detection, and foreground-window queries all
    operate on the real user desktop instead of Session 0. Reads the effective
    capture policy published by the SYSTEM check-in process and publishes live
    telemetry back through the shared ProgramData files.

    Supports screenshot mode (existing) and H.264 video segment mode (DVI-730).
    In video mode: captures at capture_fps for _SEGMENT_DURATION_S seconds,
    encodes via PyAV libx264 (in-process, SAC-safe), and sends via scp.
    Falls back to screenshot mode when PyAV is not installed.
    """
    log = logging.getLogger(__name__)
    log.info("Capture-only worker started (machine_id=%s, session=%s, frames_dir=%s, av=%s)",
             machine_id, _current_session_id(), frames_dir, _HAS_AV)
    log.info("paramiko=%s scp=%s ssh=%s key=%s key_exists=%s endpoint=%s",
             _HAS_PARAMIKO, _SCP_EXE, _SSH_EXE,
             _RSYNC_SSH_KEY, _RSYNC_SSH_KEY.exists(), _rsync_endpoint())
    _worker_token = _load_agent_token()
    _last_log_upload = 0.0
    last_ts = None
    while True:
        pol = _read_shared_policy() or dict(_cap_policy)
        mode = pol.get("capture_mode", "screenshot")
        cap_enabled = pol.get("capture_enabled", False)
        idle = _get_idle_seconds()
        fapp = _get_foreground_app()

        # On-demand screenshot (admin collect_screenshot command) always uses
        # screenshot path regardless of capture_mode.  Clear the flag first.
        screenshot_requested = pol.pop("screenshot_requested", False)
        if screenshot_requested:
            _write_shared_policy(pol)
            ts_shot = _capture_frame(machine_id, frames_dir)
            if ts_shot is not None:
                last_ts = ts_shot
                machine_dir = frames_dir / machine_id
                _send_frames(machine_dir, machine_id, _frame_files(machine_dir, ts_shot))

        if mode == "video" and _HAS_AV and cap_enabled:
            if _should_capture(pol, idle, fapp):
                fps = max(0.1, min(30.0, float(pol.get("capture_fps", 1.0))))
                new_ts = _capture_video_segment(machine_id, frames_dir, pol, fps)
                if new_ts is not None:
                    last_ts = new_ts
            else:
                time.sleep(5)  # brief poll while gated
            _write_shared_state(idle, fapp, last_ts)
            # No extra sleep — _capture_video_segment provides ~10 s of pacing.
        else:
            # Screenshot mode (or video fallback when PyAV unavailable).
            rate = max(10, pol.get("capture_rate_s", 60))
            if cap_enabled and _should_capture(pol, idle, fapp):
                ts_shot = _capture_frame(machine_id, frames_dir)
                if ts_shot is not None:
                    last_ts = ts_shot
                    machine_dir = frames_dir / machine_id
                    push_ok = _send_frames(machine_dir, machine_id, _frame_files(machine_dir, ts_shot))
                    if not push_ok and (time.time() - _last_log_upload) > 600:
                        _worker_upload_capture_log(base_url, _worker_token)
                        _last_log_upload = time.time()
            _write_shared_state(idle, fapp, last_ts)
            time.sleep(rate)


# ---------------------------------------------------------------------------
# Enrollment
# ---------------------------------------------------------------------------

def enroll(base_url: str, enrollment_token: str, machine_id: str, hostname: str) -> tuple[str, str]:
    """POST /vision/agents/register and persist the returned agent_token.

    Returns (agent_token, status) where status is 'active' or 'pending'.
    Admin-issued machine-bound tokens → 'active'.
    NULL-machine admin tokens or the open-enroll shared secret → 'pending'
    (admin must approve in the Enroll tab before capture begins).

    Raises SystemExit(1) on any network or server-side error so install.ps1
    sees a clean non-zero exit code rather than an unformatted Python traceback.
    """
    log = logging.getLogger(__name__)
    url = f"{base_url.rstrip('/')}/vision/agents/register"
    payload = {
        "enrollment_token": enrollment_token,
        "machine_id": machine_id,
        "hostname": hostname,
    }
    log.info("Enrolling with %s (machine_id=%s)", url, machine_id)
    # Use separate connect (10 s) and read (30 s) timeouts.  A single timeout=30
    # can stall much longer on Windows when a firewall silently drops the SYN
    # (TCP connect retransmit timer may exceed the socket timeout in that case).
    try:
        resp = requests.post(url, json=payload, timeout=(10, 30))
    except requests.exceptions.ConnectTimeout:
        log.error(
            "ENROLLMENT FAILED — connection to %s timed out (10 s). "
            "Check that this machine can reach the Togen server: "
            "run  Test-NetConnection %s -Port 443  in PowerShell.",
            base_url, base_url.split("//")[-1].split("/")[0],
        )
        sys.exit(1)
    except requests.exceptions.ConnectionError as exc:
        log.error(
            "ENROLLMENT FAILED — could not connect to %s: %s. "
            "Check network/firewall and DNS resolution for the Togen server.",
            base_url, exc,
        )
        sys.exit(1)
    except requests.exceptions.Timeout as exc:
        log.error(
            "ENROLLMENT FAILED — request to %s timed out: %s. "
            "The server may be unreachable or overloaded.",
            url, exc,
        )
        sys.exit(1)
    if not resp.ok:
        try:
            err_body = resp.json().get("error", resp.text[:200])
        except Exception:
            err_body = resp.text[:200]
        log.error(
            "ENROLLMENT FAILED — server returned HTTP %d: %s",
            resp.status_code, err_body,
        )
        sys.exit(1)
    data = resp.json()
    agent_token = data["agent_token"]
    status = data.get("status", "active")
    _store_agent_token(agent_token)
    if status == "pending":
        log.info(
            "Enrolled as PENDING (machine_id=%s). "
            "An admin must approve this agent in the Enroll tab before "
            "capture begins. The agent will poll and activate automatically.",
            machine_id,
        )
    else:
        log.info("Enrolled as ACTIVE — agent_token stored in credential manager.")
    return agent_token, status


# ---------------------------------------------------------------------------
# Win32 data collection
# ---------------------------------------------------------------------------

def _get_foreground_app() -> str | None:
    """Title of the foreground window; None if locked or no foreground window."""
    try:
        user32 = ctypes.windll.user32
        hwnd = user32.GetForegroundWindow()
        if not hwnd:
            return None
        length = user32.GetWindowTextLengthW(hwnd) + 1
        buf = ctypes.create_unicode_buffer(length)
        user32.GetWindowTextW(hwnd, buf, length)
        text = buf.value.strip()
        return text or None
    except Exception:
        return None


def _get_idle_seconds() -> int:
    """Seconds since last keyboard/mouse input via GetLastInputInfo."""
    class _LASTINPUTINFO(ctypes.Structure):
        _fields_ = [("cbSize", ctypes.c_uint), ("dwTime", ctypes.c_uint)]

    lii = _LASTINPUTINFO()
    lii.cbSize = ctypes.sizeof(_LASTINPUTINFO)
    if ctypes.windll.user32.GetLastInputInfo(ctypes.byref(lii)):
        elapsed_ms = ctypes.windll.kernel32.GetTickCount() - lii.dwTime
        return max(0, elapsed_ms // 1000)
    return 0


def _get_screen_count() -> int:
    """Number of attached monitors via EnumDisplayMonitors."""
    _count = [0]

    @ctypes.WINFUNCTYPE(
        ctypes.c_bool,
        ctypes.c_void_p,
        ctypes.c_void_p,
        ctypes.POINTER(ctypes.wintypes.RECT),
        ctypes.c_long,
    )
    def _cb(hm, hdc, lprect, lparam):
        _count[0] += 1
        return True

    ctypes.windll.user32.EnumDisplayMonitors(None, None, _cb, 0)
    return _count[0] or 1


def _get_logged_in_user() -> str:
    """DOMAIN\\username of the current user (mirrors VOYAGETECH\\jsmith format)."""
    domain = os.environ.get("USERDOMAIN") or platform.node()
    user = os.environ.get("USERNAME") or os.environ.get("USER") or "unknown"
    return f"{domain}\\{user}"


# ---------------------------------------------------------------------------
# Check-in payload
# ---------------------------------------------------------------------------

def _build_payload(machine_id: str, hostname: str, uptime_s: int) -> dict:
    with _pol_lock:
        cap_enabled = _cap_policy.get("capture_enabled", False)
        local_last_ts = _cap_state.get("last_ts")
    # Prefer telemetry from the interactive capture worker (DVI-683). The SYSTEM
    # check-in process runs in Session 0, where idle/foreground queries return
    # meaningless always-idle / no-window values; the in-session worker publishes
    # the real desktop state. Fall back to local values when no worker is active.
    shared = _read_shared_state()
    if shared:
        idle_seconds = shared.get("idle_seconds")
        foreground_app = shared.get("foreground_app")
        last_ts = shared.get("last_capture_ts")
    else:
        idle_seconds = _get_idle_seconds()
        foreground_app = _get_foreground_app()
        last_ts = local_last_ts
    return {
        "machine_id": machine_id,
        "hostname": hostname,
        "logged_in_user": _get_logged_in_user(),
        "os_version": platform.version(),
        "agent_version": _AGENT_VERSION,
        "status": "active",
        "capture_enabled": cap_enabled,
        "last_capture_ts": last_ts,
        "foreground_app": foreground_app,
        "idle_seconds": idle_seconds,
        "cpu_pct": psutil.cpu_percent(interval=0.5),
        "screen_count": _get_screen_count(),
        "agent_uptime_s": uptime_s,
        "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
    }


# ---------------------------------------------------------------------------
# Check-in loop
# ---------------------------------------------------------------------------

def run_checkin_loop(base_url: str, agent_token: str, machine_id: str,
                     hostname: str, frames_dir: Path, initial_pending: bool = False):
    """Main check-in and capture loop.

    When initial_pending=True (self-enrolled, awaiting admin approval), the capture
    daemon is deferred until the server reports agent_status='active' in a check-in
    response.  Check-ins still run so the server can record the agent fingerprint and
    the admin can approve from the Enroll tab.
    """
    log = logging.getLogger(__name__)
    url = f"{base_url.rstrip('/')}/vision/checkin"
    start = time.monotonic()
    interval = _DEFAULT_INTERVAL_S
    _capture_started = False
    # In Session 0 (the SYSTEM service) screen capture cannot see the interactive
    # desktop, so the dedicated AtLogOn capture worker (--capture-only) owns it and
    # this loop only publishes policy + forwards the worker's telemetry (DVI-683).
    in_session_zero = _is_session_zero()

    # DVI-804: SYSTEM creates id_togen_vision with SYSTEM-only ACL; the capture
    # worker (logged-in user) cannot read it.  Use /reset+/inheritance:r+/grant:r
    # to clear ssh-keygen DENY ACEs and set final permissions.  Runs every check-in
    # cycle as SYSTEM so the ACL stays correct even after key regeneration.
    if sys.platform == "win32" and _RSYNC_SSH_KEY.exists():
        try:
            _kp = str(_RSYNC_SSH_KEY)
            subprocess.run(["icacls", _kp, "/reset"],        capture_output=True, timeout=15)
            subprocess.run(["icacls", _kp, "/inheritance:r"],capture_output=True, timeout=15)
            subprocess.run(["icacls", _kp, "/grant:r", "NT AUTHORITY\\SYSTEM:(F)"],
                           capture_output=True, timeout=15)
            r = subprocess.run(["icacls", _kp, "/grant:r", "*S-1-5-11:(R)"],
                               capture_output=True, timeout=15)
            log.info("icacls: key ACL fixed (rc=%d) %s", r.returncode, _kp)
        except Exception as _icacls_exc:
            log.warning("icacls fix failed (non-fatal): %s", _icacls_exc)

    def _start_capture_if_needed():
        nonlocal _capture_started
        if _capture_started or in_session_zero:
            return
        _capture_started = True
        capture_t = threading.Thread(
            target=_run_capture_thread,
            args=(machine_id, frames_dir),
            daemon=True,
        )
        capture_t.start()
        log.info("Capture daemon started (in-process, interactive session).")

    if in_session_zero:
        log.info(
            "Running in Session 0 (SYSTEM). Screen capture is delegated to the "
            "interactive AtLogOn capture worker; this process handles check-in, "
            "policy publication, and telemetry forwarding only."
        )
        # Mark started so the check-in loop doesn't spam "Agent approved — starting
        # capture daemon" every cycle; the AtLogOn task owns capture in Session 0.
        _capture_started = True
    elif not initial_pending:
        _start_capture_if_needed()
    else:
        log.info(
            "Agent is PENDING approval. Check-ins running; capture suppressed "
            "until an admin approves this agent in the Enroll tab."
        )

    while True:
        uptime_s = int(time.monotonic() - start)
        payload = _build_payload(machine_id, hostname, uptime_s)

        try:
            resp = requests.post(
                url,
                json=payload,
                headers={"Authorization": f"Bearer {agent_token}"},
                timeout=15,
            )
            if resp.status_code == 401:
                log.error("agent_token rejected (401). Re-run with --enroll <token>.")
                sys.exit(1)
            resp.raise_for_status()
            data = resp.json()

            # Persist successful check-in time for the 30-day self-destruct timer.
            _set_last_checkin_at(time.time())

            # Watch for pending→active/revoked transitions (DVI-438 / DVI-735).
            server_status = data.get("agent_status")
            if server_status == "revoked":
                # Agent has been revoked by an admin; self-uninstall.
                # Ack the uninstall command if one was delivered alongside the revoke signal.
                if data.get("command") and data["command"].get("action") == "uninstall":
                    cmd_id_revoke = data["command"].get("id")
                    if cmd_id_revoke:
                        try:
                            requests.post(
                                f"{base_url.rstrip('/')}/vision/commands/{cmd_id_revoke}/ack",
                                json={"status": "ok"},
                                headers={"Authorization": f"Bearer {agent_token}"},
                                timeout=10,
                            )
                        except Exception:
                            pass
                log.warning("Server reports agent revoked — initiating self-uninstall.")
                _self_uninstall("revoked by Togen admin")
            if server_status == "active":
                if not _capture_started:
                    log.info("Agent approved — starting capture daemon.")
                _start_capture_if_needed()
            elif server_status == "pending" and not _capture_started:
                log.debug("Agent still pending approval — capture suppressed.")

            # Apply server-issued capture policy (capture_enabled, rate, exclusions…).
            if "config_delta" in data:
                _apply_config_delta(data["config_delta"])
                # Publish the effective policy so the interactive capture worker
                # (separate process / session) picks it up without its own check-in.
                if isinstance(data["config_delta"], dict):
                    _write_shared_policy(data["config_delta"])
                log.debug("Capture policy updated from config_delta")

            # Respect server-issued polling interval.
            if "next_interval_s" in data:
                new_interval = int(data["next_interval_s"])
                if new_interval != interval:
                    log.info("Check-in interval updated: %ds → %ds", interval, new_interval)
                    interval = new_interval

            # Command dispatch — ack every command so the server dequeues it.
            # Without acking, the server re-delivers the same command on every
            # check-in indefinitely (see /vision/commands/<cmd_id>/ack).
            if data.get("command"):
                cmd = data["command"]
                cmd_id = cmd.get("id")
                action = cmd.get("action", "")
                params = cmd.get("params") or {}

                ack_status = "ok"
                try:
                    if action == "pause_capture":
                        with _pol_lock:
                            _cap_policy["capture_enabled"] = False
                        _patch_shared_policy({"capture_enabled": False})
                        log.info("Command: capture paused")
                    elif action == "resume_capture":
                        with _pol_lock:
                            _cap_policy["capture_enabled"] = True
                        _patch_shared_policy({"capture_enabled": True})
                        _start_capture_if_needed()
                        log.info("Command: capture resumed")
                    elif action == "set_interval":
                        new_iv = int(params.get("interval_s", interval))
                        if new_iv != interval:
                            log.info("Command: interval set %ds → %ds", interval, new_iv)
                            interval = new_iv
                    elif action == "collect_screenshot":
                        # SYSTEM (Session 0) cannot screenshot the desktop; signal the
                        # interactive capture worker via the shared policy file.
                        _patch_shared_policy({"screenshot_requested": True})
                        log.info("Command: screenshot requested (forwarded to capture worker)")
                    elif action == "restart_agent":
                        log.info("Command: restart requested — exiting (Scheduled Task will restart)")
                        if cmd_id:
                            try:
                                requests.post(
                                    f"{base_url.rstrip('/')}/vision/commands/{cmd_id}/ack",
                                    json={"status": "ok"},
                                    headers={"Authorization": f"Bearer {agent_token}"},
                                    timeout=10,
                                )
                            except Exception:
                                pass
                        sys.exit(0)
                    elif action == "uninstall":
                        log.info("Command: uninstall requested — initiating self-destruct")
                        # Ack before uninstalling; _self_uninstall calls sys.exit.
                        if cmd_id:
                            try:
                                requests.post(
                                    f"{base_url.rstrip('/')}/vision/commands/{cmd_id}/ack",
                                    json={"status": "ok"},
                                    headers={"Authorization": f"Bearer {agent_token}"},
                                    timeout=10,
                                )
                            except Exception:
                                pass
                        _self_uninstall("uninstall command received from Togen")
                    elif action == "start_capture_task":
                        result = subprocess.run(
                            ["schtasks", "/run", "/tn", "TogenVisionAgentCapture"],
                            capture_output=True,
                            creationflags=_NO_WINDOW,
                        )
                        if result.returncode == 0:
                            log.info("Command: capture task started")
                        else:
                            log.warning("Command: schtasks /run failed (rc=%d): %s",
                                        result.returncode, result.stderr.decode(errors="replace").strip())
                            ack_status = "error"
                    elif action == "update_agent":
                        import shutil as _shutil
                        import tempfile as _tempfile
                        # DVI-850: Python handles ALL networking (requests works through
                        # corporate proxy; SYSTEM+PowerShell Invoke-WebRequest often
                        # fails).  We download both upgrade.ps1 and the agent ZIP here,
                        # then pass the ZIP path to upgrade.ps1 via -PreDownloadedZip so
                        # the script needs zero network access.
                        # Script saved to install dir (not %TEMP%) for WDAC compatibility.
                        # utf-8-sig (BOM) so PowerShell 5.1 reads encoding correctly.
                        # CREATE_BREAKAWAY_FROM_JOB so upgrade survives Stop-ScheduledTask.
                        _srv = base_url.rstrip("/")
                        log.info("Command: update_agent — downloading upgrade.ps1 + ZIP")
                        try:
                            # 1. Download upgrade.ps1
                            _ps1_resp = requests.get(
                                f"{_srv}/enroll/upgrade.ps1",
                                headers={"Authorization": f"Bearer {agent_token}"},
                                timeout=60,
                            )
                            _ps1_resp.raise_for_status()

                            # 2. Download agent ZIP (~49 MB) — Python handles proxy
                            log.info("Command: downloading agent ZIP from %s/enroll/agent.zip", _srv)
                            _zip_resp = requests.get(
                                f"{_srv}/enroll/agent.zip",
                                headers={"Authorization": f"Bearer {agent_token}"},
                                timeout=600,
                                stream=True,
                            )
                            _zip_resp.raise_for_status()
                            _zip_tmp = _tempfile.NamedTemporaryFile(
                                prefix="togen-agent-", suffix=".zip", delete=False
                            )
                            _zip_bytes = 0
                            try:
                                for _chunk in _zip_resp.iter_content(chunk_size=1 << 20):
                                    _zip_tmp.write(_chunk)
                                    _zip_bytes += len(_chunk)
                            finally:
                                _zip_tmp.close()
                            _zip_path = _zip_tmp.name
                            log.info("Command: agent ZIP downloaded (%d bytes) to %s", _zip_bytes, _zip_path)
                            if _zip_bytes < 10 * 1024 * 1024:
                                raise ValueError(f"ZIP too small: {_zip_bytes} bytes")

                            # 3. Write upgrade.ps1 to install dir
                            _install_dir = os.path.join(
                                os.environ.get("ProgramFiles", r"C:\Program Files"),
                                "Togen", "Vision",
                            )
                            os.makedirs(_install_dir, exist_ok=True)
                            ps1 = os.path.join(_install_dir, "upgrade.ps1")
                            with open(ps1, "w", encoding="utf-8-sig") as _fz:
                                _fz.write(_ps1_resp.text)

                            ps_exe = (
                                _shutil.which("powershell.exe")
                                or r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe"
                            )

                            # 4. Post Python-level diagnostic so we can see upgrade started
                            try:
                                requests.post(
                                    f"{_srv}/vision/agent-log",
                                    json={"file": "upgrade_diag.txt",
                                          "content": (
                                              f"py_upgrade_start=true ps1={ps1} "
                                              f"zip_size={_zip_bytes} zip_path={_zip_path} "
                                              f"ps_exe={ps_exe}"
                                          )},
                                    headers={"Authorization": f"Bearer {agent_token}"},
                                    timeout=10,
                                )
                            except Exception:
                                pass

                            # 5. Launch upgrade.ps1 with pre-downloaded ZIP path
                            _CF_BREAKAWAY = 0x01000000  # CREATE_BREAKAWAY_FROM_JOB
                            try:
                                subprocess.Popen(
                                    [ps_exe, "-ExecutionPolicy", "Bypass",
                                     "-File", ps1, "-PreDownloadedZip", _zip_path],
                                    cwd=_install_dir,
                                    creationflags=subprocess.CREATE_NO_WINDOW
                                    | subprocess.CREATE_NEW_PROCESS_GROUP
                                    | _CF_BREAKAWAY,
                                )
                            except OSError:
                                subprocess.Popen(
                                    [ps_exe, "-ExecutionPolicy", "Bypass",
                                     "-File", ps1, "-PreDownloadedZip", _zip_path],
                                    cwd=_install_dir,
                                    creationflags=subprocess.CREATE_NO_WINDOW
                                    | subprocess.CREATE_NEW_PROCESS_GROUP,
                                )
                            log.info("Command: upgrade launched (ps1=%s zip=%s)", ps1, _zip_path)
                        except Exception as exc:
                            log.warning("Command: update_agent failed: %s", exc)
                            ack_status = "error"
                    elif action == "upload_log":
                        # Flush and upload all on-disk log files to the server log
                        # store so the admin log viewer (DVI-705 / DVI-709) can read them.
                        _LOG_MAX_UPLOAD = 200_000  # server caps at 200k anyway
                        _log_files = [_AGENT_LOG_FILE, _CAPTURE_LOG_FILE]
                        uploaded = 0
                        for _lf in _log_files:
                            try:
                                if not _lf.exists():
                                    continue
                                _text = _lf.read_text(encoding="utf-8", errors="replace")
                                if len(_text) > _LOG_MAX_UPLOAD:
                                    _text = _text[-_LOG_MAX_UPLOAD:]
                                requests.post(
                                    f"{base_url.rstrip('/')}/vision/agent-log",
                                    json={"file": _lf.name, "content": _text},
                                    headers={"Authorization": f"Bearer {agent_token}"},
                                    timeout=30,
                                )
                                uploaded += 1
                                log.info("Command: uploaded log %s (%d chars)", _lf.name, len(_text))
                            except Exception as _ul_exc:
                                log.warning("Command: upload_log %s failed: %s", _lf.name, _ul_exc)
                                ack_status = "error"
                        if uploaded == 0 and ack_status == "ok":
                            log.info("Command: upload_log — no log files found at %s", _SHARE_DIR)
                    elif action == "cleanup_legacy":
                        legacy_removed = _cleanup_legacy_vision()
                        log.info("Command: cleanup_legacy done — %d item(s) removed: %s",
                                 len(legacy_removed), legacy_removed or "none")
                    else:
                        log.info("Command %s not handled by this agent version", action)
                        ack_status = "error"
                except Exception as exc:
                    log.warning("Error handling command %s: %s", action, exc)
                    ack_status = "error"

                if cmd_id:
                    try:
                        requests.post(
                            f"{base_url.rstrip('/')}/vision/commands/{cmd_id}/ack",
                            json={"status": ack_status},
                            headers={"Authorization": f"Bearer {agent_token}"},
                            timeout=10,
                        )
                    except Exception as exc:
                        log.warning("Failed to ack command %s: %s", cmd_id, exc)

            log.debug("Check-in OK (uptime=%ds, interval=%ds)", uptime_s, interval)

        except requests.RequestException as exc:
            log.warning("Check-in failed: %s", exc)

        time.sleep(interval)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Togen Vision Windows Agent")
    parser.add_argument(
        "--server", required=True,
        help="Base URL of the Togen server (e.g. https://togen.example.com)",
    )
    parser.add_argument(
        "--enroll", metavar="ENROLLMENT_TOKEN",
        help=(
            "Enroll this machine. Accepts either an admin-issued short-lived token "
            "(quick-enroll → active) or the open-enroll shared secret from the admin "
            "panel (self-enroll → pending, admin must approve in the Enroll tab)."
        ),
    )
    parser.add_argument(
        "--enroll-only", action="store_true",
        help=(
            "Enroll and exit without starting the check-in loop. "
            "Used by install.ps1; the Scheduled Task handles the ongoing run."
        ),
    )
    parser.add_argument(
        "--capture-only", action="store_true",
        help=(
            "Run only the screen-capture worker (no check-in). Launched by the "
            "AtLogOn scheduled task in the interactive user session so capture, "
            "idle detection, and foreground-window queries see the real desktop. "
            "Reads the effective policy published by the SYSTEM check-in process."
        ),
    )
    parser.add_argument("--debug", action="store_true", help="Enable debug logging.")
    parser.add_argument(
        "--frames-dir", metavar="PATH",
        help="Local directory for captured frames (default: ~/.togen/frames). "
             "Override via TOGEN_FRAMES_DIR env var.",
    )
    args = parser.parse_args()

    if args.enroll_only and not args.enroll:
        parser.error("--enroll-only requires --enroll <token>")

    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO,
        format=_LOG_FMT,
    )
    log = logging.getLogger(__name__)

    # Rotate on-disk log so the admin log viewer (DVI-705 / DVI-709) can fetch it.
    # The capture worker writes to its own file to avoid the Windows file-lock
    # contention that previously dropped its lines from agent.log (DVI-798).
    _log_target = _CAPTURE_LOG_FILE if args.capture_only else _AGENT_LOG_FILE
    try:
        _SHARE_DIR.mkdir(parents=True, exist_ok=True)
        _fh = logging.handlers.RotatingFileHandler(
            _log_target,
            maxBytes=_AGENT_LOG_MAX_BYTES,
            backupCount=_AGENT_LOG_BACKUP_COUNT,
            encoding="utf-8",
        )
        _fh.setFormatter(logging.Formatter(_LOG_FMT))
        logging.getLogger().addHandler(_fh)
        log.info("File logging started: %s", _log_target)
    except Exception as _fh_err:
        log.warning("Could not open log file %s: %s", _log_target, _fh_err)

    machine_id = get_machine_id()
    hostname = socket.gethostname()
    frames_dir = Path(args.frames_dir) if args.frames_dir else _FRAMES_DIR

    # Remove old C:\Local\ PowerShell-era Vision deployments (DVI-722).
    # Runs on every startup — safe, idempotent, fast when nothing is present.
    if not args.capture_only:
        _cleanup_legacy_vision()

    if args.capture_only:
        if not _HAS_PIL:
            log.error("Pillow not installed — capture worker cannot run. "
                      "Install it: pip install Pillow")
            sys.exit(1)
        run_capture_only(machine_id, frames_dir, args.server)
        return

    initial_pending = False
    if args.enroll:
        _, status = enroll(args.server, args.enroll, machine_id, hostname)
        initial_pending = (status == "pending")
        if args.enroll_only:
            # Exit code 0 = active, 2 = pending (install.ps1 uses this to print guidance).
            sys.exit(2 if initial_pending else 0)

    agent_token = _load_agent_token()
    if not agent_token:
        log.error("No agent_token found. Enroll first: --enroll <token>")
        sys.exit(1)

    if not _HAS_PIL:
        log.warning(
            "Pillow not installed — screen capture disabled. "
            "Install it: pip install Pillow"
        )

    # Check whether the agent has been offline for >30 days and should self-destruct.
    _check_self_destruct_timeout()

    log.info("Vision agent starting (machine_id=%s, hostname=%s, frames_dir=%s)",
             machine_id, hostname, frames_dir)
    run_checkin_loop(args.server, agent_token, machine_id, hostname, frames_dir,
                     initial_pending=initial_pending)


if __name__ == "__main__":
    main()
