"""Guest command transport over vsock-compatible JSON protocol."""

from __future__ import annotations

import base64
import json
import socket
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Protocol


class SocketLike(Protocol):
    def settimeout(self, timeout: int) -> None: ...

    def connect(self, address: Any) -> None: ...

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

    def recv(self, size: int) -> bytes: ...

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


SocketFactory = Callable[[int, int], SocketLike]


@dataclass(frozen=True)
class GuestExecResponse:
    stdout: str
    stderr: str
    exit_code: int
    duration_ms: int


@dataclass(frozen=True)
class GuestArchiveResponse:
    destination: str
    entry_count: int
    bytes_written: int


@dataclass(frozen=True)
class GuestArchiveExportResponse:
    workspace_path: str
    artifact_type: str
    entry_count: int
    bytes_written: int


@dataclass(frozen=True)
class GuestWorkspaceFileReadResponse:
    path: str
    size_bytes: int
    content_bytes: bytes


@dataclass(frozen=True)
class GuestShellSummary:
    shell_id: str
    cwd: str
    cols: int
    rows: int
    state: str
    started_at: float
    ended_at: float | None
    exit_code: int | None


@dataclass(frozen=True)
class GuestShellReadResponse(GuestShellSummary):
    cursor: int
    next_cursor: int
    output: str
    truncated: bool


class VsockExecClient:
    """Minimal JSON-over-stream client for a guest exec agent."""

    def __init__(self, socket_factory: SocketFactory | None = None) -> None:
        self._socket_factory = socket_factory or socket.socket

    def exec(
        self,
        guest_cid: int,
        port: int,
        command: str,
        timeout_seconds: int,
        *,
        env: dict[str, str] | None = None,
        uds_path: str | None = None,
    ) -> GuestExecResponse:
        payload = self._request_json(
            guest_cid,
            port,
            {
                "command": command,
                "timeout_seconds": timeout_seconds,
                "env": env,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest exec response must be a JSON object",
        )
        return GuestExecResponse(
            stdout=str(payload.get("stdout", "")),
            stderr=str(payload.get("stderr", "")),
            exit_code=int(payload.get("exit_code", -1)),
            duration_ms=int(payload.get("duration_ms", 0)),
        )

    def upload_archive(
        self,
        guest_cid: int,
        port: int,
        archive_path: Path,
        *,
        destination: str,
        timeout_seconds: int = 60,
        uds_path: str | None = None,
    ) -> GuestArchiveResponse:
        request = {
            "action": "extract_archive",
            "destination": destination,
            "archive_size": archive_path.stat().st_size,
        }
        sock = self._connect(guest_cid, port, timeout_seconds, uds_path=uds_path)
        try:
            sock.sendall((json.dumps(request) + "\n").encode("utf-8"))
            with archive_path.open("rb") as handle:
                for chunk in iter(lambda: handle.read(65536), b""):
                    sock.sendall(chunk)
            payload = self._recv_json_payload(sock)
        finally:
            sock.close()

        if not isinstance(payload, dict):
            raise RuntimeError("guest archive response must be a JSON object")
        error = payload.get("error")
        if error is not None:
            raise RuntimeError(str(error))
        return GuestArchiveResponse(
            destination=str(payload.get("destination", destination)),
            entry_count=int(payload.get("entry_count", 0)),
            bytes_written=int(payload.get("bytes_written", 0)),
        )

    def install_secrets(
        self,
        guest_cid: int,
        port: int,
        archive_path: Path,
        *,
        timeout_seconds: int = 60,
        uds_path: str | None = None,
    ) -> GuestArchiveResponse:
        request = {
            "action": "install_secrets",
            "archive_size": archive_path.stat().st_size,
        }
        sock = self._connect(guest_cid, port, timeout_seconds, uds_path=uds_path)
        try:
            sock.sendall((json.dumps(request) + "\n").encode("utf-8"))
            with archive_path.open("rb") as handle:
                for chunk in iter(lambda: handle.read(65536), b""):
                    sock.sendall(chunk)
            payload = self._recv_json_payload(sock)
        finally:
            sock.close()

        if not isinstance(payload, dict):
            raise RuntimeError("guest secret install response must be a JSON object")
        error = payload.get("error")
        if error is not None:
            raise RuntimeError(str(error))
        return GuestArchiveResponse(
            destination=str(payload.get("destination", "/run/pyro-secrets")),
            entry_count=int(payload.get("entry_count", 0)),
            bytes_written=int(payload.get("bytes_written", 0)),
        )

    def export_archive(
        self,
        guest_cid: int,
        port: int,
        *,
        workspace_path: str,
        archive_path: Path,
        timeout_seconds: int = 60,
        uds_path: str | None = None,
    ) -> GuestArchiveExportResponse:
        request = {
            "action": "export_archive",
            "path": workspace_path,
        }
        sock = self._connect(guest_cid, port, timeout_seconds, uds_path=uds_path)
        try:
            sock.sendall((json.dumps(request) + "\n").encode("utf-8"))
            header = self._recv_line(sock)
            if header.strip() == "":
                raise RuntimeError("guest export response header is empty")
            payload = json.loads(header)
            if not isinstance(payload, dict):
                raise RuntimeError("guest export response header must be a JSON object")
            error = payload.get("error")
            if error is not None:
                raise RuntimeError(str(error))
            archive_size = int(payload.get("archive_size", 0))
            if archive_size < 0:
                raise RuntimeError("guest export archive_size must not be negative")
            with archive_path.open("wb") as handle:
                remaining = archive_size
                while remaining > 0:
                    chunk = sock.recv(min(65536, remaining))
                    if chunk == b"":
                        raise RuntimeError("unexpected EOF while receiving export archive")
                    handle.write(chunk)
                    remaining -= len(chunk)
        finally:
            sock.close()
        return GuestArchiveExportResponse(
            workspace_path=str(payload.get("workspace_path", workspace_path)),
            artifact_type=str(payload.get("artifact_type", "file")),
            entry_count=int(payload.get("entry_count", 0)),
            bytes_written=int(payload.get("bytes_written", 0)),
        )

    def list_workspace_entries(
        self,
        guest_cid: int,
        port: int,
        *,
        workspace_path: str,
        recursive: bool,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        return self._request_json(
            guest_cid,
            port,
            {
                "action": "list_workspace",
                "path": workspace_path,
                "recursive": recursive,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest workspace file list response must be a JSON object",
        )

    def read_workspace_file(
        self,
        guest_cid: int,
        port: int,
        *,
        workspace_path: str,
        max_bytes: int,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        payload = self._request_json(
            guest_cid,
            port,
            {
                "action": "read_workspace_file",
                "path": workspace_path,
                "max_bytes": max_bytes,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest workspace file read response must be a JSON object",
        )
        raw_content = payload.get("content_b64", "")
        if not isinstance(raw_content, str):
            raise RuntimeError("guest workspace file read response is missing content_b64")
        payload["content_bytes"] = base64.b64decode(raw_content.encode("ascii"), validate=True)
        payload.pop("content_b64", None)
        return payload

    def write_workspace_file(
        self,
        guest_cid: int,
        port: int,
        *,
        workspace_path: str,
        text: str,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        return self._request_json(
            guest_cid,
            port,
            {
                "action": "write_workspace_file",
                "path": workspace_path,
                "text": text,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest workspace file write response must be a JSON object",
        )

    def delete_workspace_path(
        self,
        guest_cid: int,
        port: int,
        *,
        workspace_path: str,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        return self._request_json(
            guest_cid,
            port,
            {
                "action": "delete_workspace_path",
                "path": workspace_path,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest workspace path delete response must be a JSON object",
        )

    def open_shell(
        self,
        guest_cid: int,
        port: int,
        *,
        shell_id: str,
        cwd: str,
        cols: int,
        rows: int,
        env: dict[str, str] | None = None,
        redact_values: list[str] | None = None,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> GuestShellSummary:
        payload = self._request_json(
            guest_cid,
            port,
            {
                "action": "open_shell",
                "shell_id": shell_id,
                "cwd": cwd,
                "cols": cols,
                "rows": rows,
                "env": env,
                "redact_values": redact_values,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest shell open response must be a JSON object",
        )
        return self._shell_summary_from_payload(payload)

    def read_shell(
        self,
        guest_cid: int,
        port: int,
        *,
        shell_id: str,
        cursor: int,
        max_chars: int,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> GuestShellReadResponse:
        payload = self._request_json(
            guest_cid,
            port,
            {
                "action": "read_shell",
                "shell_id": shell_id,
                "cursor": cursor,
                "max_chars": max_chars,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest shell read response must be a JSON object",
        )
        summary = self._shell_summary_from_payload(payload)
        return GuestShellReadResponse(
            shell_id=summary.shell_id,
            cwd=summary.cwd,
            cols=summary.cols,
            rows=summary.rows,
            state=summary.state,
            started_at=summary.started_at,
            ended_at=summary.ended_at,
            exit_code=summary.exit_code,
            cursor=int(payload.get("cursor", cursor)),
            next_cursor=int(payload.get("next_cursor", cursor)),
            output=str(payload.get("output", "")),
            truncated=bool(payload.get("truncated", False)),
        )

    def write_shell(
        self,
        guest_cid: int,
        port: int,
        *,
        shell_id: str,
        input_text: str,
        append_newline: bool,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        payload = self._request_json(
            guest_cid,
            port,
            {
                "action": "write_shell",
                "shell_id": shell_id,
                "input": input_text,
                "append_newline": append_newline,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest shell write response must be a JSON object",
        )
        self._shell_summary_from_payload(payload)
        return payload

    def signal_shell(
        self,
        guest_cid: int,
        port: int,
        *,
        shell_id: str,
        signal_name: str,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        payload = self._request_json(
            guest_cid,
            port,
            {
                "action": "signal_shell",
                "shell_id": shell_id,
                "signal": signal_name,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest shell signal response must be a JSON object",
        )
        self._shell_summary_from_payload(payload)
        return payload

    def close_shell(
        self,
        guest_cid: int,
        port: int,
        *,
        shell_id: str,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        payload = self._request_json(
            guest_cid,
            port,
            {
                "action": "close_shell",
                "shell_id": shell_id,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest shell close response must be a JSON object",
        )
        self._shell_summary_from_payload(payload)
        return payload

    def start_service(
        self,
        guest_cid: int,
        port: int,
        *,
        service_name: str,
        command: str,
        cwd: str,
        readiness: dict[str, Any] | None,
        ready_timeout_seconds: int,
        ready_interval_ms: int,
        env: dict[str, str] | None = None,
        timeout_seconds: int = 60,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        return self._request_json(
            guest_cid,
            port,
            {
                "action": "start_service",
                "service_name": service_name,
                "command": command,
                "cwd": cwd,
                "readiness": readiness,
                "ready_timeout_seconds": ready_timeout_seconds,
                "ready_interval_ms": ready_interval_ms,
                "env": env,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest service start response must be a JSON object",
        )

    def status_service(
        self,
        guest_cid: int,
        port: int,
        *,
        service_name: str,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        return self._request_json(
            guest_cid,
            port,
            {
                "action": "status_service",
                "service_name": service_name,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest service status response must be a JSON object",
        )

    def logs_service(
        self,
        guest_cid: int,
        port: int,
        *,
        service_name: str,
        tail_lines: int | None,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        return self._request_json(
            guest_cid,
            port,
            {
                "action": "logs_service",
                "service_name": service_name,
                "tail_lines": tail_lines,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest service logs response must be a JSON object",
        )

    def stop_service(
        self,
        guest_cid: int,
        port: int,
        *,
        service_name: str,
        timeout_seconds: int = 30,
        uds_path: str | None = None,
    ) -> dict[str, Any]:
        return self._request_json(
            guest_cid,
            port,
            {
                "action": "stop_service",
                "service_name": service_name,
            },
            timeout_seconds=timeout_seconds,
            uds_path=uds_path,
            error_message="guest service stop response must be a JSON object",
        )

    def _request_json(
        self,
        guest_cid: int,
        port: int,
        request: dict[str, Any],
        *,
        timeout_seconds: int,
        uds_path: str | None,
        error_message: str,
    ) -> dict[str, Any]:
        sock = self._connect(guest_cid, port, timeout_seconds, uds_path=uds_path)
        try:
            sock.sendall((json.dumps(request) + "\n").encode("utf-8"))
            payload = self._recv_json_payload(sock)
        finally:
            sock.close()
        if not isinstance(payload, dict):
            raise RuntimeError(error_message)
        error = payload.get("error")
        if error is not None:
            raise RuntimeError(str(error))
        return payload

    @staticmethod
    def _shell_summary_from_payload(payload: dict[str, Any]) -> GuestShellSummary:
        return GuestShellSummary(
            shell_id=str(payload.get("shell_id", "")),
            cwd=str(payload.get("cwd", "/workspace")),
            cols=int(payload.get("cols", 0)),
            rows=int(payload.get("rows", 0)),
            state=str(payload.get("state", "stopped")),
            started_at=float(payload.get("started_at", 0.0)),
            ended_at=(
                None if payload.get("ended_at") is None else float(payload.get("ended_at", 0.0))
            ),
            exit_code=(
                None if payload.get("exit_code") is None else int(payload.get("exit_code", 0))
            ),
        )

    def _connect(
        self,
        guest_cid: int,
        port: int,
        timeout_seconds: int,
        *,
        uds_path: str | None,
    ) -> SocketLike:
        family = getattr(socket, "AF_VSOCK", None)
        if family is not None:
            sock = self._socket_factory(family, socket.SOCK_STREAM)
            connect_address: Any = (guest_cid, port)
        elif uds_path is not None:
            sock = self._socket_factory(socket.AF_UNIX, socket.SOCK_STREAM)
            connect_address = uds_path
        else:
            raise RuntimeError("vsock sockets are not supported on this host Python runtime")
        sock.settimeout(timeout_seconds)
        sock.connect(connect_address)
        if family is None:
            sock.sendall(f"CONNECT {port}\n".encode("utf-8"))
            status = self._recv_line(sock)
            if not status.startswith("OK "):
                sock.close()
                raise RuntimeError(f"vsock unix bridge rejected port {port}: {status.strip()}")
        return sock

    @staticmethod
    def _recv_line(sock: SocketLike) -> str:
        chunks: list[bytes] = []
        while True:
            data = sock.recv(1)
            if data == b"":
                break
            chunks.append(data)
            if data == b"\n":
                break
        return b"".join(chunks).decode("utf-8", errors="replace")

    @staticmethod
    def _recv_json_payload(sock: SocketLike) -> Any:
        chunks: list[bytes] = []
        while True:
            data = sock.recv(65536)
            if data == b"":
                break
            chunks.append(data)
        return json.loads(b"".join(chunks).decode("utf-8"))
