"""FastVM implementation of OSWorld's `VMManager`.

FastVM — like AWS/Azure/GCP — is a cloud-style allocator: VMs are created
on demand and torn down at the end of an episode, so the registry methods
(`add_vm`/`delete_vm`/`occupy_vm`/…) are no-ops. The only meaningful work
is in `get_vm_path`, which POSTs `/v1/vms` and blocks until the VM is
running (or fails).

The returned `path_to_vm` is the FastVM VM id (e.g. `"85fd32bb-…"`);
the provider uses it as the subject for every subsequent SDK call.
"""
from __future__ import annotations

import logging
import os
import signal
import time
from typing import Any, Optional

from desktop_env.providers.base import VMManager
from desktop_env.providers.fastvm.client import get_client
from desktop_env.providers.fastvm.config import (
    DEFAULT_MACHINE_TYPE,
    DISK_GIB,
    LAUNCH_TIMEOUT_SEC,
    OSWORLD_SNAPSHOT_ID,
)

logger = logging.getLogger("desktopenv.providers.fastvm.FastvmVMManager")
logger.setLevel(logging.INFO)


def _osworld_firewall() -> dict:
    """Keep the policy in one place so provider + manager agree."""
    return {"mode": "open"}


def _allocate_vm(snapshot_id: Optional[str], machine_type: str) -> str:
    """Launch a new VM and wait for `running`.

    Installs SIGINT/SIGTERM handlers that tear down a partially-launched VM
    so Ctrl-C during the long `provisioning` → `running` transition doesn't
    leak a VM. Mirrors `aws.manager._allocate_vm`.
    """
    client = get_client()
    vm_id: Optional[str] = None
    prev_int = signal.getsignal(signal.SIGINT)
    prev_term = signal.getsignal(signal.SIGTERM)

    def _handler(sig, _frame):
        if vm_id:
            name = "SIGINT" if sig == signal.SIGINT else "SIGTERM"
            logger.warning(f"{name} received, deleting FastVM {vm_id}…")
            try:
                client.vms.delete(vm_id)
            except Exception as cleanup_err:  # noqa: BLE001
                logger.error(f"Cleanup after {name} failed: {cleanup_err}")
        signal.signal(signal.SIGINT, prev_int)
        signal.signal(signal.SIGTERM, prev_term)
        if sig == signal.SIGINT:
            raise KeyboardInterrupt
        raise SystemExit(0)

    try:
        signal.signal(signal.SIGINT, _handler)
        signal.signal(signal.SIGTERM, _handler)

        launch_kwargs: dict[str, Any] = {
            "firewall": _osworld_firewall(),
            "name": f"osworld-{int(time.time())}",
            "wait_timeout": LAUNCH_TIMEOUT_SEC,
        }
        if snapshot_id:
            launch_kwargs["snapshot_id"] = snapshot_id
            # machine_type is inherited from the snapshot source.
        else:
            launch_kwargs["machine_type"] = machine_type
        if DISK_GIB:
            launch_kwargs["disk_gi_b"] = DISK_GIB

        logger.info(
            "Launching FastVM (snapshot_id=%s, machine_type=%s)…",
            snapshot_id,
            machine_type if not snapshot_id else "inherited",
        )
        vm = client.vms.launch(**launch_kwargs)
        vm_id = vm.id
        logger.info(f"FastVM {vm_id} is {vm.status} (public_ipv6={vm.public_ipv6})")
        return vm_id

    except Exception:
        if vm_id:
            logger.info(f"Tearing down FastVM {vm_id} after error")
            try:
                get_client().vms.delete(vm_id)
            except Exception as cleanup_err:  # noqa: BLE001
                logger.error(f"Cleanup after error failed: {cleanup_err}")
        raise
    finally:
        signal.signal(signal.SIGINT, prev_int)
        signal.signal(signal.SIGTERM, prev_term)


class FastvmVMManager(VMManager):
    """FastVM has no persistent VM pool — everything is allocate-on-demand."""

    def __init__(self, **_kwargs: Any) -> None:
        self.initialize_registry()

    def initialize_registry(self, **_kwargs: Any) -> None:
        pass

    def add_vm(self, vm_path, **_kwargs):
        pass

    def delete_vm(self, vm_path, **_kwargs):
        pass

    def occupy_vm(self, vm_path, pid, **_kwargs):
        pass

    def list_free_vms(self, **_kwargs):
        return []

    def check_and_clean(self, **_kwargs):
        pass

    def get_vm_path(self, **kwargs: Any) -> str:
        """Allocate a VM for one OSWorld episode.

        `os_type`/`region`/`screen_size` are accepted for interface parity but
        ignored — FastVM's image catalogue is selected via a snapshot id,
        and screen resolution lives inside the guest image (xorg conf).
        """
        snapshot_id = kwargs.get("snapshot_id") or OSWORLD_SNAPSHOT_ID
        machine_type = kwargs.get("machine_type") or DEFAULT_MACHINE_TYPE
        if not snapshot_id:
            logger.warning(
                "No FASTVM_OSWORLD_SNAPSHOT_ID set — booting from the default "
                "FastVM image. The OSWorld in-VM server won't be present, so "
                "the readiness probe will fail unless FASTVM_SKIP_READY_CHECK=1."
            )
        return _allocate_vm(snapshot_id=snapshot_id, machine_type=machine_type)
