from __future__ import annotations

import socket

import pytest

from pyro_mcp.vm_guest import VsockExecClient


class StubSocket:
    def __init__(self, response: bytes) -> None:
        self.response = response
        self.connected: tuple[int, int] | None = None
        self.sent = b""
        self.timeout: int | None = None
        self.closed = False

    def settimeout(self, timeout: int) -> None:
        self.timeout = timeout

    def connect(self, address: tuple[int, int]) -> None:
        self.connected = address

    def sendall(self, data: bytes) -> None:
        self.sent += data

    def recv(self, size: int) -> bytes:
        del size
        if self.response == b"":
            return b""
        data, self.response = self.response, b""
        return data

    def close(self) -> None:
        self.closed = True


def test_vsock_exec_client_round_trip(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(socket, "AF_VSOCK", 40, raising=False)
    stub = StubSocket(
        b'{"stdout":"ok\\n","stderr":"","exit_code":0,"duration_ms":7}'
    )

    def socket_factory(family: int, sock_type: int) -> StubSocket:
        assert family == socket.AF_VSOCK
        assert sock_type == socket.SOCK_STREAM
        return stub

    client = VsockExecClient(socket_factory=socket_factory)
    response = client.exec(1234, 5005, "echo ok", 30)

    assert response.exit_code == 0
    assert response.stdout == "ok\n"
    assert stub.connected == (1234, 5005)
    assert b'"command": "echo ok"' in stub.sent
    assert stub.closed is True


def test_vsock_exec_client_rejects_bad_json(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(socket, "AF_VSOCK", 40, raising=False)
    stub = StubSocket(b"[]")
    client = VsockExecClient(socket_factory=lambda family, sock_type: stub)
    with pytest.raises(RuntimeError, match="JSON object"):
        client.exec(1234, 5005, "echo ok", 30)