"""Smoke test for the FastVM provider against api.fastvm.org.

Runs end-to-end: launch → query IP → snapshot → revert → delete. Kept
isolated from `DesktopEnv` so we don't need the full OSWorld ML stack
(torch, transformers, …) to verify the control-plane integration.

Exits 0 on success, non-zero on any failure. Prints a timing summary
so we can eyeball the restore-from-snapshot latency.

Usage:
    FASTVM_API_KEY=... python -m desktop_env.providers.fastvm.smoke_test
    # Add FASTVM_OSWORLD_SNAPSHOT_ID=... to exercise the snapshot path.
    # Set FASTVM_SKIP_READY_CHECK=1 to skip the in-VM :5000 probe (for
    # images that don't run osworld_server).
"""
from __future__ import annotations

import logging
import os
import sys
import time

from desktop_env.providers.fastvm.manager import FastvmVMManager
from desktop_env.providers.fastvm.provider import FastvmProvider

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(name)s %(levelname)s %(message)s",
)


def _fmt(seconds: float) -> str:
    return f"{seconds:.1f}s"


def main() -> int:
    if not os.environ.get("FASTVM_API_KEY"):
        print("ERROR: FASTVM_API_KEY not set", file=sys.stderr)
        return 2

    timings: dict[str, float] = {}

    manager = FastvmVMManager()
    provider = FastvmProvider()

    # Collect snapshots we create so we can clean them up at the end; the
    # provider intentionally doesn't auto-delete snapshots (callers may want
    # them for warm starts), so bookkeeping belongs here in the smoke test.
    created_snapshots: list[str] = []

    t0 = time.monotonic()
    vm_id = manager.get_vm_path()
    timings["launch"] = time.monotonic() - t0
    print(f"[ok] launched vm_id={vm_id} in {_fmt(timings['launch'])}")

    try:
        t0 = time.monotonic()
        ip_ports = provider.get_ip_address(vm_id)
        timings["get_ip"] = time.monotonic() - t0
        print(f"[ok] ip_ports={ip_ports} in {_fmt(timings['get_ip'])}")

        # Snapshot/revert is independent of the OSWorld image — FastVM snapshots
        # a running VM regardless of what's installed in it. The readiness probe
        # only guards the :5000 check; keep snapshot/revert always-on unless
        # explicitly gated with FASTVM_SMOKE_SKIP_SNAPSHOT=1.
        if os.environ.get("FASTVM_SMOKE_SKIP_SNAPSHOT") == "1":
            print("[skip] snapshot/revert (FASTVM_SMOKE_SKIP_SNAPSHOT=1)")
            return 0

        t0 = time.monotonic()
        snap_id = provider.save_state(vm_id, f"osworld-smoke-{int(time.time())}")
        timings["snapshot"] = time.monotonic() - t0
        created_snapshots.append(snap_id)
        print(f"[ok] snapshot={snap_id} in {_fmt(timings['snapshot'])}")

        t0 = time.monotonic()
        vm_id = provider.revert_to_snapshot(vm_id, snap_id)
        timings["revert"] = time.monotonic() - t0
        print(f"[ok] reverted to vm_id={vm_id} in {_fmt(timings['revert'])}")

        t0 = time.monotonic()
        ip_ports2 = provider.get_ip_address(vm_id)
        timings["get_ip_restored"] = time.monotonic() - t0
        print(f"[ok] restored ip_ports={ip_ports2} in {_fmt(timings['get_ip_restored'])}")

    finally:
        print("[teardown] deleting VM")
        try:
            provider.stop_emulator(vm_id)
        except Exception as e:  # noqa: BLE001
            print(f"[warn] teardown failed: {e}")
        for sid in created_snapshots:
            try:
                provider._client.snapshots.delete(sid)  # type: ignore[attr-defined]
                print(f"[teardown] deleted snapshot {sid}")
            except Exception as e:  # noqa: BLE001
                print(f"[warn] snapshot teardown failed for {sid}: {e}")

    print("\n=== Timings ===")
    for k, v in timings.items():
        print(f"  {k:<20} {_fmt(v)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
