"""Power and snapshot operations (the write half of the HypervisorDriver contract)."""

from __future__ import annotations

from typing import Any

from pyVmomi import vim

from napalm_vmware._tasks import invoke, wait_for_task, wait_until
from napalm_vmware.parse.snapshots import find_snapshot, snapshot_list

_TASK_TIMEOUT = 300
_GUEST_SHUTDOWN_TIMEOUT = 300
_RUNNING = "poweredOn"
_TOOLS_RUNNING = "guestToolsRunning"


class VmwareActionsMixin:
    """Mixed into both drivers ahead of :class:`VmwareBaseDriver`."""

    _inventory: Any
    _find_vm: Any
    _mo: Any

    def _vm_mo(self, vm: dict[str, Any]) -> Any:
        return self._mo(vim.VirtualMachine, vm["_moref"])

    def _run(self, call: Any, *args: Any, **kwargs: Any) -> None:
        wait_for_task(self._inventory, invoke(call, *args, **kwargs), _TASK_TIMEOUT)

    @staticmethod
    def _require_running(vm: dict[str, Any], action: str) -> None:
        if vm.get("runtime.powerState") != _RUNNING:
            raise RuntimeError(f"Cannot {action} VM {vm.get('name')!r}: it is not running")

    @staticmethod
    def _require_tools(vm: dict[str, Any], action: str) -> None:
        if vm.get("guest.toolsRunningStatus") != _TOOLS_RUNNING:
            raise RuntimeError(
                f"A graceful {action} of {vm.get('name')!r} needs VMware Tools running "
                "in the guest; use force to cut power instead"
            )

    # -- power ---------------------------------------------------------------

    def start_vm(self, name: str) -> None:
        self._run(self._vm_mo(self._find_vm(name)).PowerOnVM_Task)

    def stop_vm(self, name: str, force: bool = False) -> None:
        vm = self._find_vm(name)
        self._require_running(vm, "stop")
        mo = self._vm_mo(vm)
        if force:
            self._run(mo.PowerOffVM_Task)
            return
        self._require_tools(vm, "shutdown")
        invoke(mo.ShutdownGuest)
        wait_until(
            lambda: (
                self._inventory.properties(mo, ["runtime.powerState"]).get("runtime.powerState")
                == "poweredOff"
            ),
            _GUEST_SHUTDOWN_TIMEOUT,
            f"guest shutdown of {vm.get('name')!r}",
        )

    def reboot_vm(self, name: str, force: bool = False) -> None:
        vm = self._find_vm(name)
        self._require_running(vm, "reboot")
        mo = self._vm_mo(vm)
        if force:
            self._run(mo.ResetVM_Task)
            return
        # RebootGuest returns once Tools accepted the request; the guest gives
        # no signal when it is back, so there is nothing further to wait for.
        self._require_tools(vm, "reboot")
        invoke(mo.RebootGuest)

    def suspend_vm(self, name: str) -> None:
        vm = self._find_vm(name)
        self._require_running(vm, "suspend")
        self._run(self._vm_mo(vm).SuspendVM_Task)

    # -- snapshots -------------------------------------------------------------

    def create_vm_snapshot(
        self, name: str, snapshot: str, description: str = "", include_memory: bool = False
    ) -> None:
        vm = self._find_vm(name)
        if any(s["name"] == snapshot for s in snapshot_list(vm)):
            raise ValueError(f"VM {vm.get('name')!r} already has a snapshot named {snapshot!r}")
        self._run(
            self._vm_mo(vm).CreateSnapshot_Task,
            name=snapshot,
            description=description,
            memory=include_memory and vm.get("runtime.powerState") == _RUNNING,
            quiesce=False,
        )

    def _snapshot_mo(self, name: str, snapshot: str) -> Any:
        return self._mo(vim.vm.Snapshot, find_snapshot(self._find_vm(name), snapshot))

    def delete_vm_snapshot(self, name: str, snapshot: str) -> None:
        # removeChildren=False consolidates the snapshot into its children
        # instead of discarding them.
        self._run(self._snapshot_mo(name, snapshot).RemoveSnapshot_Task, removeChildren=False)

    def rollback_vm_snapshot(self, name: str, snapshot: str) -> None:
        self._run(self._snapshot_mo(name, snapshot).RevertToSnapshot_Task)