"""Bake an OSWorld-ready FastVM snapshot.

One-shot: launch a VM → upload the OSWorld server payload + install script →
run the install → verify `:5000/screenshot` responds from outside the VM →
snapshot → print the snapshot id → delete the VM.

Usage:
    # One-time setup, run once per FastVM organisation:
    FASTVM_API_KEY=... python scripts/fastvm/bake_image.py \
        --mode minimal \
        --name osworld-minimal-v1

    # Full desktop image (LibreOffice/Chrome/VLC/…) — ~30–45 min:
    FASTVM_API_KEY=... python scripts/fastvm/bake_image.py \
        --mode full \
        --machine-type c4m8 \
        --disk-gib 40 \
        --name osworld-full-v1

Stdout on success ends with the exported-env-var line the operator can
source directly:

    export FASTVM_OSWORLD_SNAPSHOT_ID=<uuid>

We do NOT delete the VM on failure so the operator can SSH in (via console)
and diagnose. Pass --cleanup-on-error to override.
"""
from __future__ import annotations

import argparse
import logging
import os
import pathlib
import shlex
import subprocess
import sys
import time
from typing import Optional

from fastvm import FastvmClient
from fastvm.lib import VMLaunchError, VMNotReadyError

REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
SERVER_SRC = REPO_ROOT / "desktop_env" / "server"
SCRIPTS_DIR = pathlib.Path(__file__).resolve().parent

# Remote layout inside the VM. Kept in one place so the install scripts,
# systemd units, and this bake script all agree.
REMOTE_SERVER_DIR = "/opt/osworld-server"
REMOTE_INSTALL_SCRIPT = {
    "minimal": f"{REMOTE_SERVER_DIR}/install_osworld_minimal.sh",
    "full": f"{REMOTE_SERVER_DIR}/install_osworld_full.sh",
}

# Install-step timeouts — minimal is a short apt-get + pip install; full
# pulls ~1 GB of apps and rebuilds fontconfig, which is slow.
INSTALL_TIMEOUT_SEC = {
    "minimal": 10 * 60,
    "full": 60 * 60,
}

LOG = logging.getLogger("bake_image")


def _upload_payload(client: FastvmClient, vm_id: str, mode: str) -> None:
    """Copy the OSWorld server sources + install scripts into the VM.

    We upload the whole `desktop_env/server/` directory (main.py,
    pyxcursor.py, requirements.txt, osworld_server.service) plus both
    install scripts so they're available for future re-bakes without
    re-uploading.
    """
    LOG.info("Uploading server payload to %s:%s", vm_id, REMOTE_SERVER_DIR)
    # `client.upload(local_dir, remote_dir)` tars the dir and extracts
    # server-side — one round trip.
    client.upload(vm_id, str(SERVER_SRC), REMOTE_SERVER_DIR)
    # The install scripts live in scripts/fastvm/ locally but we stash them
    # alongside the payload so the full installer can shell out to the
    # minimal one via an absolute path.
    for mode_key in ("minimal", "full"):
        local = SCRIPTS_DIR / f"install_osworld_{mode_key}.sh"
        remote = REMOTE_INSTALL_SCRIPT[mode_key]
        client.upload(vm_id, str(local), remote)
    # Execute bit — `client.upload` preserves perms from the local file if
    # unpacked via tar (directory upload) but the per-file path uses a GCS
    # presign that strips them. Chmod defensively.
    client.vms.run(
        vm_id,
        command=f"chmod +x {shlex.quote(REMOTE_INSTALL_SCRIPT['minimal'])} "
        f"{shlex.quote(REMOTE_INSTALL_SCRIPT['full'])}",
        timeout_sec=60,
    )


def _run_install(client: FastvmClient, vm_id: str, mode: str) -> None:
    """Fire-and-forget the install script, then poll for its exit code.

    The full install can take 20+ min. A single long `client.vms.run` call
    trips the SDK's ~60 s HTTP read timeout, the SDK retries, and the
    install ends up running twice — which is not idempotent for
    unit-register/start steps. Instead: launch via `nohup` and poll a
    sentinel file. Each exec round-trip is sub-second.
    """
    timeout = INSTALL_TIMEOUT_SEC[mode]
    LOG.info("Running %s install (poll-timeout %ds)", mode, timeout)
    script = REMOTE_INSTALL_SCRIPT[mode]
    log_path = "/tmp/osworld-install.log"
    rc_path = "/tmp/osworld-install.rc"

    # Write a launcher on the VM that backgrounds the install and records
    # the exit code. Doing this as a guest-side file (rather than a single
    # shell string with embedded `&` / `setsid`) sidesteps the guest-agent
    # quoting quirks that made earlier attempts fail with
    # "guest exec transport error".
    launcher_remote = "/opt/osworld-server/_bake_launcher.sh"
    # Launcher design: one nohup'd child runs the whole recipe — install,
    # then write exit code. The parent launcher exits as soon as it fires
    # the child, so the exec API returns in <1 s. Wrapping both steps in
    # ONE bash -c keeps them in the same process (no orphaned `wait`s,
    # no disown/job-control tricks that previously made $? race to 127).
    launcher_body = (
        "#!/bin/sh\n"
        f"rm -f {shlex.quote(rc_path)} {shlex.quote(log_path)}\n"
        f"nohup bash -c "
        f"{shlex.quote(f'bash -ex {shlex.quote(script)} > {shlex.quote(log_path)} 2>&1; echo $? > {shlex.quote(rc_path)}')}"
        f" > /dev/null 2>&1 &\n"
        "echo launched\n"
    )
    import tempfile as _tempfile
    with _tempfile.NamedTemporaryFile("w", suffix=".sh", delete=False) as f:
        f.write(launcher_body)
        launcher_local = f.name
    try:
        client.upload(vm_id, launcher_local, launcher_remote)
    finally:
        os.unlink(launcher_local)
    client.vms.run(vm_id, command=f"chmod +x {shlex.quote(launcher_remote)}", timeout_sec=30)

    kick = client.vms.run(vm_id, command=launcher_remote, timeout_sec=30)
    if kick.exit_code != 0 or "launched" not in (kick.stdout or ""):
        raise SystemExit(
            f"failed to launch install script: rc={kick.exit_code} "
            f"stdout={kick.stdout!r} stderr={kick.stderr!r}"
        )

    # Poll the rc sentinel. Start with a short interval so a fast minimal
    # install doesn't waste 30s of idle polling.
    deadline = time.monotonic() + timeout
    interval = 5.0
    while time.monotonic() < deadline:
        probe = client.vms.run(
            vm_id,
            command=f"cat {shlex.quote(rc_path)} 2>/dev/null || echo RUNNING",
            timeout_sec=15,
        )
        rc_text = (probe.stdout or "").strip()
        if rc_text and rc_text != "RUNNING":
            try:
                exit_code = int(rc_text)
            except ValueError:
                raise SystemExit(f"install rc file contained unexpected value: {rc_text!r}")
            if exit_code != 0:
                LOG.error("Install script failed (exit=%s)", exit_code)
                tail = client.vms.run(
                    vm_id,
                    command=f"tail -n 200 {shlex.quote(log_path)}",
                    timeout_sec=30,
                )
                sys.stderr.write("--- install log tail ---\n")
                sys.stderr.write((tail.stdout or "")[-8000:])
                sys.stderr.write("\n")
                raise SystemExit(f"install failed on VM {vm_id}")
            LOG.info("Install script completed.")
            return
        LOG.info("Install still running…")
        time.sleep(interval)
        # Geometric backoff up to 30s; quick check early, lazy later.
        interval = min(interval * 1.3, 30.0)

    raise SystemExit(f"install did not finish within {timeout}s on VM {vm_id}")


def _verify_from_outside(ipv6: str, timeout: float = 120.0) -> None:
    """Hit the VM's `:5000/screenshot` over the public network.

    Proving the endpoint works from *outside* the VM — i.e. through the
    FastVM firewall, not just via localhost inside the guest — is the
    whole point of a "ready" snapshot.
    """
    url = f"http://[{ipv6}]:5000/screenshot"
    LOG.info("Probing %s", url)
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        # Use curl rather than requests so we don't re-learn v6-bracket URL
        # parsing the hard way on older httpx versions.
        p = subprocess.run(
            ["curl", "-6", "--max-time", "5", "-sS", "-o", "/dev/null", "-w", "%{http_code}", url],
            capture_output=True,
            text=True,
        )
        if p.returncode == 0 and p.stdout.strip() == "200":
            LOG.info("External probe OK (HTTP 200)")
            return
        LOG.info("Probe not ready yet (curl_rc=%s http=%s)", p.returncode, p.stdout.strip() or "?")
        time.sleep(3)
    raise SystemExit(f"External :5000 probe never returned 200 within {timeout:.0f}s")


def _create_and_wait_snapshot(client: FastvmClient, vm_id: str, name: str, timeout: float = 600.0) -> str:
    LOG.info("Creating snapshot %r of %s", name, vm_id)
    snap = client.snapshots.create(vm_id=vm_id, name=name)
    deadline = time.monotonic() + timeout
    while True:
        current = next((s for s in client.snapshots.list() if s.id == snap.id), None)
        if current is None:
            raise SystemExit(f"snapshot {snap.id} disappeared mid-poll")
        if current.status == "ready":
            LOG.info("Snapshot %s is ready.", snap.id)
            return snap.id
        if current.status == "error":
            raise SystemExit(f"snapshot {snap.id} entered error state")
        if time.monotonic() > deadline:
            raise SystemExit(f"snapshot {snap.id} not ready within {timeout:.0f}s")
        time.sleep(2)


def _launch(client: FastvmClient, *, machine_type: str, disk_gib: Optional[int]) -> str:
    kwargs = {
        "machine_type": machine_type,
        "firewall": {"mode": "open"},
        "name": f"osworld-bake-{int(time.time())}",
        "wait_timeout": 300.0,
    }
    if disk_gib:
        kwargs["disk_gi_b"] = disk_gib
    LOG.info("Launching bake VM (%s)", machine_type)
    vm = client.vms.launch(**kwargs)
    LOG.info("Bake VM %s is %s (ipv6=%s)", vm.id, vm.status, vm.public_ipv6)
    return vm.id


def main(argv: Optional[list[str]] = None) -> int:
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument("--mode", choices=("minimal", "full"), default="minimal",
                   help="minimal: Xvfb + server only (~5 min). full: desktop + apps (~30–45 min).")
    p.add_argument("--name", default=None,
                   help="Snapshot name. Defaults to osworld-<mode>-<timestamp>.")
    p.add_argument("--machine-type", default="c2m4",
                   help="FastVM machine type for the bake VM. Full mode needs >= c4m8.")
    p.add_argument("--disk-gib", type=int, default=None,
                   help="Override disk size (GiB). Full mode typically wants 40.")
    p.add_argument("--cleanup-on-error", action="store_true",
                   help="Delete the bake VM if install/probe/snapshot fails. Default: leave it for debugging.")
    p.add_argument("--keep-vm", action="store_true",
                   help="Don't delete the VM after a successful bake (useful for iterating on the install script).")
    p.add_argument("-v", "--verbose", action="store_true")
    args = p.parse_args(argv)

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s %(message)s",
    )

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

    name = args.name or f"osworld-{args.mode}-{int(time.time())}"
    client = FastvmClient()

    timings: dict[str, float] = {}
    vm_id: Optional[str] = None
    start = time.monotonic()

    try:
        t = time.monotonic()
        vm_id = _launch(client, machine_type=args.machine_type, disk_gib=args.disk_gib)
        timings["launch"] = time.monotonic() - t

        # Re-fetch VM for the IPv6 (launch() returns with it populated but be
        # defensive in case we retry).
        vm = client.vms.retrieve(vm_id)
        if not vm.public_ipv6:
            raise SystemExit(f"VM {vm_id} has no publicIpv6")

        t = time.monotonic()
        _upload_payload(client, vm_id, args.mode)
        timings["upload"] = time.monotonic() - t

        t = time.monotonic()
        _run_install(client, vm_id, args.mode)
        timings["install"] = time.monotonic() - t

        t = time.monotonic()
        _verify_from_outside(vm.public_ipv6)
        timings["verify_external"] = time.monotonic() - t

        t = time.monotonic()
        snapshot_id = _create_and_wait_snapshot(client, vm_id, name)
        timings["snapshot"] = time.monotonic() - t

    except (VMLaunchError, VMNotReadyError, SystemExit):
        if vm_id and args.cleanup_on_error:
            LOG.warning("Cleaning up bake VM %s on error", vm_id)
            try:
                client.vms.delete(vm_id)
            except Exception as e:  # noqa: BLE001
                LOG.error("Cleanup failed: %s", e)
        elif vm_id:
            LOG.warning("Leaving bake VM %s in place for debugging (pass --cleanup-on-error to delete)", vm_id)
        raise

    timings["total"] = time.monotonic() - start

    # Success — tear down unless --keep-vm.
    if vm_id and not args.keep_vm:
        try:
            client.vms.delete(vm_id)
        except Exception as e:  # noqa: BLE001
            LOG.warning("Post-bake VM delete failed (harmless): %s", e)

    print("\n=== Bake complete ===")
    for k, v in timings.items():
        print(f"  {k:<20} {v:.1f}s")
    print()
    print(f"snapshot_id: {snapshot_id}")
    print(f"snapshot_name: {name}")
    print()
    print("To use this image as the OSWorld base:")
    print(f"    export FASTVM_OSWORLD_SNAPSHOT_ID={snapshot_id}")
    print()
    return 0


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