"""Host-side, one-shot guest execution over Firecracker's vsock Unix socket."""

from __future__ import annotations

import socket
import stat
import struct
import time
from dataclasses import dataclass
from pathlib import Path

from .errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError
from .guest_agent import (
    AGENT_PORT, CLEANUP_SECONDS, CONNECT_TIMEOUT_SECONDS, MAX_OUTPUT_BYTES,
    MAX_REQUEST_BYTES, MAX_RESPONSE_BYTES, PROTOCOL_VERSION, ProtocolError,
    receive_exact, receive_frame, send_frame, validate_command,
)


@dataclass(frozen=True, slots=True)
class ExecResult:
    exit_code: int | None
    stdout: str
    stderr: str
    timed_out: bool
    truncated: bool


def validate_exec(command: object, sudo: object, timeout_seconds: object) -> None:
    try:
        validate_command(command, sudo, timeout_seconds)
    except ProtocolError as error:
        raise ValidationError(str(error)) from error


class GuestSession:
    """A connected session pins the selected VMM; never reconnect after submission."""

    def __init__(self, connection: socket.socket) -> None:
        self._connection = connection

    @classmethod
    def connect(cls, path: Path, vm_id: str, username: str, pid: int) -> GuestSession:
        connection = None
        try:
            connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            metadata = path.lstat()
            if not stat.S_ISSOCK(metadata.st_mode):
                raise ProtocolError("guest socket is not a Unix socket")
            deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
            connection.settimeout(CONNECT_TIMEOUT_SECONDS)
            connection.connect(str(path))
            peer_pid, _uid, _gid = struct.unpack("3i", connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
            if peer_pid != pid:
                raise ProtocolError("guest socket does not belong to the selected VM process")
            connection.sendall(f"CONNECT {AGENT_PORT}\n".encode("ascii"))
            acknowledgement = bytearray()
            while not acknowledgement.endswith(b"\n") and len(acknowledgement) < 32:
                acknowledgement.extend(receive_exact(connection, 1, deadline))
            parts = bytes(acknowledgement).split()
            if (not acknowledgement.endswith(b"\n") or len(parts) != 2
                    or parts[0] != b"OK" or not parts[1].isdigit()):
                raise ProtocolError("invalid Firecracker vsock acknowledgement")
            hello = receive_frame(connection, 4096, deadline)
            if (set(hello) != {"version", "vm_id", "username"}
                    or type(hello["version"]) is not int or hello["version"] != PROTOCOL_VERSION
                    or hello["vm_id"] != vm_id or hello["username"] != username):
                raise ProtocolError("guest agent identity or protocol version does not match this VM")
            return cls(connection)
        except (OSError, EOFError, ProtocolError) as error:
            if connection is not None:
                connection.close()
            raise GuestAgentUnavailableError(
                "guest agent is unavailable or not ready; install it with POST /vms/{vm_id}/agent/install "
                "while stopped, then start the VM (vsock-capable kernel required)"
            ) from error

    def execute(self, command: str, *, sudo: bool, timeout_seconds: float) -> ExecResult:
        validate_exec(command, sudo, timeout_seconds)
        try:
            send_frame(self._connection, {
                "version": PROTOCOL_VERSION, "command": command,
                "sudo": sudo, "timeout_seconds": timeout_seconds,
            }, MAX_REQUEST_BYTES, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
            result = receive_frame(
                self._connection, MAX_RESPONSE_BYTES,
                time.monotonic() + timeout_seconds + CLEANUP_SECONDS + CONNECT_TIMEOUT_SECONDS,
            )
            if "error" in result:
                if (set(result) != {"error", "outcome_unknown"} or not isinstance(result["error"], str)
                        or len(result["error"]) > 4096 or type(result["outcome_unknown"]) is not bool):
                    raise ProtocolError("invalid guest failure response")
                result["error"].encode("utf-8")  # Reject invalid Unicode before rendering an HTTP error.
                raise GuestExecutionError(result["error"], outcome_unknown=result["outcome_unknown"])
            if set(result) != {"exit_code", "stdout", "stderr", "timed_out", "truncated"}:
                raise ProtocolError("invalid guest result fields")
            if (type(result["timed_out"]) is not bool or type(result["truncated"]) is not bool
                    or not isinstance(result["stdout"], str) or not isinstance(result["stderr"], str)
                    or (result["timed_out"] and result["exit_code"] is not None)
                    or (not result["timed_out"] and (
                        type(result["exit_code"]) is not int or not -64 <= result["exit_code"] <= 255))):
                raise ProtocolError("invalid guest result types")
            if sum(len(result[name].encode("utf-8")) for name in ("stdout", "stderr")) > 3 * MAX_OUTPUT_BYTES:
                raise ProtocolError("guest output exceeds its limit")
            return ExecResult(**result)
        except (OSError, EOFError, ProtocolError, UnicodeError) as error:
            raise GuestExecutionError(
                "guest execution connection was lost or returned an invalid result; "
                "the command may have run and will not be retried",
                outcome_unknown=True,
            ) from error

    def close(self) -> None:
        self._connection.close()