"""Interactive, process-pinned guest terminal transport over the existing vsock agent."""

from __future__ import annotations

import asyncio
import base64
import json
import struct
import time

from .errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError
from .guest_agent import (
    CONNECT_TIMEOUT_SECONDS, MAX_TERMINAL_FRAME, PROTOCOL_VERSION, ProtocolError,
    receive_frame, send_frame, terminal_data, validate_terminal_size,
)
from .guest_exec import GuestSession


def validate_size(rows: object, cols: object) -> None:
    try:
        validate_terminal_size(rows, cols)
    except ProtocolError as error:
        raise ValidationError(str(error)) from error


class GuestTerminal(GuestSession):
    def start(self, rows: int, cols: int) -> None:
        validate_size(rows, cols)
        try:
            deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
            send_frame(self._connection, {"version": PROTOCOL_VERSION, "terminal": {"rows": rows, "cols": cols}},
                       MAX_TERMINAL_FRAME, deadline)
            response = receive_frame(self._connection, MAX_TERMINAL_FRAME, deadline)
            if response != {"type": "ready"}:
                raise ProtocolError("terminal support is unavailable")
            self._connection.setblocking(False)
            self._write_lock = asyncio.Lock()
        except (OSError, EOFError, ProtocolError) as error:
            self.close()
            raise GuestAgentUnavailableError(
                "Interactive terminal support is unavailable. Restart this VM to apply automatic guest-agent setup, "
                "then wait for the agent to start and launch a new session."
            ) from error

    async def _receive_exact(self, size: int) -> bytes:
        data = bytearray()
        while len(data) < size:
            chunk = await asyncio.get_running_loop().sock_recv(self._connection, size - len(data))
            if not chunk:
                raise EOFError("guest terminal connection closed")
            data.extend(chunk)
        return bytes(data)

    async def read(self) -> dict:
        try:
            # An idle shell has no output deadline; partial messages are bounded.
            first = await self._receive_exact(1)
            async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
                size, = struct.unpack("!I", first + await self._receive_exact(3))
                if not 0 < size <= MAX_TERMINAL_FRAME:
                    raise ProtocolError("invalid terminal frame length")
                message = json.loads(await self._receive_exact(size))
            if isinstance(message, dict):
                if set(message) == {"type", "data"} and message["type"] == "output":
                    return {"type": "output", "data": terminal_data(message["data"])}
                if (set(message) == {"type", "exit_code"} and message["type"] == "exit"
                        and type(message["exit_code"]) is int and -64 <= message["exit_code"] <= 255):
                    return message
            raise ProtocolError("invalid guest terminal response")
        except (OSError, EOFError, ValueError, UnicodeError, RecursionError) as error:
            raise GuestExecutionError("The guest terminal connection ended unexpectedly.", outcome_unknown=True) from error

    async def _send(self, message: dict) -> None:
        payload = json.dumps(message, separators=(",", ":")).encode("ascii")
        if len(payload) > MAX_TERMINAL_FRAME:
            raise ValidationError("terminal input exceeds its size limit")
        async with self._write_lock:
            async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
                await asyncio.get_running_loop().sock_sendall(self._connection, struct.pack("!I", len(payload)) + payload)

    async def write(self, data: bytes) -> None:
        await self._send({"type": "input", "data": base64.b64encode(data).decode("ascii")})

    async def resize(self, rows: int, cols: int) -> None:
        validate_size(rows, cols)
        await self._send({"type": "resize", "rows": rows, "cols": cols})

    async def end(self) -> None:
        await self._send({"type": "close"})