"""The guest channel: qemu-guest-agent over virtio-serial.

WHY NOT SSH

The obvious way to run a command in a VM is sshd and a forwarded port,
and it was rejected for three reasons that all matter here:

  * a forwarded port is a LISTENING SOCKET on this host. The whole
    isolation argument is that the sandbox has no inbound surface at
    all, and `hostfwd=tcp::2222-:22` - the form everyone writes -
    binds 0.0.0.0, i.e. the LAN.
  * it needs key material, generated per session, living somewhere.
  * it stops working the moment the agent breaks the guest's networking,
    which is exactly the moment the agent most needs to be reachable.
    A machine that has to be healthy to be fixed is not self-healing.

virtio-serial has none of those properties. It is a character device on
the PCI bus; there is no address, no port and no key, and it works with
the guest's network stack down entirely.

WHAT "A NORMAL SUPERUSER BASH" MEANS

On Linux qemu-ga runs as root, so `guest-exec` of `/bin/bash -lc` is
literally a root login shell - a login shell rather than `-c`, so PATH
and the profile are what a person would see. On Windows the same agent
runs as LocalSystem, which is NOT a normal elevated session: no user
profile, different COM and GUI behaviour, and half of PowerShell's module
autoloading behaves oddly. That is why the Windows image carries a second
broker on its own virtio-serial port, running in the interactive admin
session, and this module is the Linux half.
"""

from __future__ import annotations

import base64
import json
import socket
import time
from typing import Any

# qemu-ga buffers a command's output in the guest before we read it, so
# this is a cap on what the guest can make vmd hold in memory, not just a
# tidiness limit. The agent's own tool layer trims again for context.
MAX_CAPTURE = 4_000_000


class GuestError(RuntimeError):
    """The guest agent is not answering, or refused a request."""


class GuestTimeout(GuestError):
    """A command outlived its timeout and was abandoned."""


class Guest:
    def __init__(self, path: str, timeout: float = 15.0,
                 os_name: str = "ubuntu") -> None:
        # Which machine this is, because the shell is not the same one.
        self.os_name = (os_name or "ubuntu").strip().lower()
        self.path = path
        self.timeout = timeout
        self._sock: socket.socket | None = None
        self._buf = b""

    def __enter__(self) -> "Guest":
        self.connect()
        return self

    def __exit__(self, *_exc: object) -> None:
        self.close()

    def connect(self) -> None:
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.settimeout(self.timeout)
        sock.connect(self.path)
        self._sock = sock
        self._resync()

    def close(self) -> None:
        if self._sock is not None:
            try:
                self._sock.close()
            finally:
                self._sock = None

    def _resync(self) -> None:
        """Discard whatever the last conversation left in the pipe.

        virtio-serial is a stream with no framing and no reset. If a
        previous command timed out, its reply is still queued, and the
        next request reads that instead - every answer one behind the
        question, which is a genuinely horrible bug to diagnose. The
        documented fix is a 0xFF byte, which the agent treats as "discard
        the parser state", followed by a sync token it echoes back.
        """
        assert self._sock is not None
        token = int(time.time() * 1000) & 0x7FFFFFFF
        self._sock.sendall(b"\xff" + json.dumps(
            {"execute": "guest-sync-delimited", "arguments": {"id": token}}
        ).encode() + b"\n")
        deadline = time.monotonic() + self.timeout
        self._buf = b""
        while time.monotonic() < deadline:
            try:
                chunk = self._sock.recv(65536)
            except socket.timeout:
                break
            if not chunk:
                break
            self._buf += chunk
            # The reply is preceded by its own 0xFF; everything before it
            # is the stale output we are trying to get rid of.
            if b"\xff" in self._buf:
                self._buf = self._buf.split(b"\xff", 1)[1]
            for line in self._buf.split(b"\n"):
                if not line.strip():
                    continue
                try:
                    message = json.loads(line)
                except ValueError:
                    continue
                if message.get("return") == token:
                    self._buf = b""
                    return
        raise GuestError(
            "the guest agent did not answer a sync. The machine is still "
            "booting, or qemu-guest-agent is not running inside it."
        )

    def _call(self, command: str, timeout: float | None = None, **arguments: Any) -> Any:
        assert self._sock is not None, "not connected"
        self._sock.settimeout(timeout or self.timeout)
        message: dict[str, Any] = {"execute": command}
        if arguments:
            message["arguments"] = arguments
        self._sock.sendall(json.dumps(message).encode() + b"\n")
        deadline = time.monotonic() + (timeout or self.timeout)
        while time.monotonic() < deadline:
            line, sep, rest = self._buf.partition(b"\n")
            if sep:
                self._buf = rest
                if not line.strip():
                    continue
                reply = json.loads(line)
                if "error" in reply:
                    raise GuestError(f"{command}: {reply['error'].get('desc', reply['error'])}")
                return reply.get("return")
            try:
                chunk = self._sock.recv(65536)
            except socket.timeout:
                break
            if not chunk:
                raise GuestError("the guest agent closed the channel")
            self._buf += chunk
        raise GuestTimeout(f"the guest agent did not answer '{command}' in time")

    # -- verbs ---------------------------------------------------------
    def ping(self, timeout: float = 5.0) -> bool:
        try:
            self._call("guest-ping", timeout=timeout)
            return True
        except (GuestError, OSError):
            return False

    def info(self) -> dict[str, Any]:
        return self._call("guest-info") or {}

    def exec(self, argv: list[str], timeout_s: float = 120.0,
             cwd: str | None = None, env: list[str] | None = None) -> dict[str, Any]:
        """Run a command and wait for it. Always returns all three of
        exit code, stdout and stderr - a tool that returns only the
        output it thinks you wanted is a tool that hides the reason
        something failed.
        """
        if not argv:
            raise GuestError("no command was given")
        arguments: dict[str, Any] = {
            "path": argv[0], "arg": argv[1:], "capture-output": True,
        }
        if env:
            arguments["env"] = env
        started = time.monotonic()
        pid = self._call("guest-exec", **arguments)["pid"]

        # Poll rather than block: the agent has no "wait" verb, and a
        # long command would otherwise sit inside a socket timeout with
        # no way to notice a cancellation.
        delay = 0.05
        while True:
            status = self._call("guest-exec-status", pid=pid)
            if status.get("exited"):
                out = base64.b64decode(status.get("out-data", "") or "")[:MAX_CAPTURE]
                err = base64.b64decode(status.get("err-data", "") or "")[:MAX_CAPTURE]
                code = status.get("exitcode")
                if code is None and "signal" in status:
                    # Reported as 128+signal so the caller sees one number.
                    code = 128 + int(status["signal"])
                return {
                    "exit_code": int(code if code is not None else -1),
                    "stdout": out.decode("utf-8", "replace"),
                    "stderr": err.decode("utf-8", "replace"),
                    "truncated": bool(status.get("out-truncated") or status.get("err-truncated")),
                    "seconds": round(time.monotonic() - started, 2),
                }
            if time.monotonic() - started > timeout_s:
                # There is no way to kill it from here - the agent
                # exposes no such verb - so say so rather than implying
                # the command stopped. It is still running in the guest.
                raise GuestTimeout(
                    f"the command was still running after {timeout_s:.0f}s and was "
                    "abandoned. It has NOT been stopped inside the machine."
                )
            time.sleep(delay)
            delay = min(delay * 1.5, 1.0)

    def shell(self, script: str, timeout_s: float = 120.0) -> dict[str, Any]:
        """A superuser shell on whichever operating system this is.

        THIS SENT `/bin/bash -lc` TO WINDOWS FOR AS LONG AS WINDOWS EXISTED.

        The module was written as the Linux half and the Windows half was
        left for the broker, so every command in a Windows session came
        back `guest-exec: Failed to execute child process (No such file or
        directory)`. That does not read as "wrong shell" to an agent - it
        reads as a machine that has stopped answering, so it rebooted the
        machine and waited, twice, and then did it again.

        PowerShell is what a Windows guest has. `-NoProfile` because
        qemu-ga runs as LocalSystem and a profile that is not there costs
        a second per call; `-NonInteractive` so a prompt fails instead of
        hanging until the timeout.

        WHAT THIS IS NOT: an interactive elevated session. qemu-ga is
        LocalSystem in session 0, so anything graphical started from here
        runs where nobody can see it, and per-user state is not the
        logged-in user's. Scripts and their output are exactly right;
        that is most of what this machine is for. The in-guest broker is
        still what makes the rest true.
        """
        if self.os_name == "windows":
            return self.exec(
                ["powershell", "-NoProfile", "-NonInteractive", "-Command", script],
                timeout_s=timeout_s,
            )
        return self.exec(["/bin/bash", "-lc", script], timeout_s=timeout_s)

    def write_file(self, path: str, content: bytes) -> int:
        handle = self._call("guest-file-open", path=path, mode="wb")
        try:
            written = 0
            # 512 KB of raw bytes per call; base64 makes each message a
            # third larger again, and the agent has a maximum message size.
            for start in range(0, len(content), 524288):
                chunk = content[start:start + 524288]
                result = self._call(
                    "guest-file-write", handle=handle,
                    **{"buf-b64": base64.b64encode(chunk).decode()},
                )
                written += int(result.get("count", 0))
            return written
        finally:
            self._call("guest-file-close", handle=handle)

    def read_file(self, path: str, max_bytes: int = 1_000_000) -> bytes:
        handle = self._call("guest-file-open", path=path, mode="rb")
        try:
            out = b""
            while len(out) < max_bytes:
                result = self._call("guest-file-read", handle=handle,
                                    count=min(524288, max_bytes - len(out)))
                out += base64.b64decode(result.get("buf-b64", "") or "")
                if result.get("eof"):
                    break
            return out
        finally:
            self._call("guest-file-close", handle=handle)

    def set_time(self, when_ns: int) -> None:
        """Set the guest's clock, in UTC nanoseconds, from the host's.

        EVERY SESSION, NOT ONLY RESUMED ONES.

        A machine resumed from a saved state wakes up believing it is
        whenever the state was taken - measured at three minutes stale on
        the first restore, and it would be days on a snapshot left alone
        over a weekend. Windows minds more than Linux: a clock that is
        wrong fails TLS certificate validation, so `pip`, `winget` and
        anything that downloads break in a way that reads as a network
        fault rather than a clock.

        The time is passed explicitly rather than letting the agent sync
        from the RTC. The RTC here is `base=localtime` for Windows and
        `base=utc` for Linux, and a verb whose meaning depends on which
        machine it is talking to is a verb that will be wrong on one of
        them.
        """
        self._call("guest-set-time", time=int(when_ns))

    def info_time(self) -> int:
        """The guest's own clock, in nanoseconds UTC. For checking a set."""
        return int(self._call("guest-get-time"))

    def shutdown(self, mode: str = "powerdown") -> None:
        if mode not in ("powerdown", "reboot", "halt"):
            raise GuestError(f"'{mode}' is not a shutdown mode")
        try:
            self._call("guest-shutdown", mode=mode, timeout=5.0)
        except (GuestError, OSError):
            # The agent goes away mid-shutdown by definition; an error
            # here means it worked, not that it failed.
            pass